Simulation & Cartesian position controller

simulation.py

@@ -0,0 +1,127 @@
+import time
+
+import numpy as np
+import pybullet as p
+import pybullet_data
+
+from sensor_msgs.msg import JointState
+
+from utils import *
+
+
+class PandaArm(object):
+    def __init__(self):
+        self.num_dof = 7
+        self.ee_link_id = 7
+        self.ee_frame_id = "panda_link7"
+
+        self.named_joint_values = {"ready": [0.0, -0.79, 0.0, -2.356, 0.0, 1.57, 0.79]}
+
+        self.upper_limits = [-7] * self.num_dof
+        self.lower_limits = [7] * self.num_dof
+        self.ranges = [7] * self.num_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 in range(self.num_dof):
+            p.resetJointState(self.uid, i, self.named_joint_values["ready"][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
+            )
+
+    def pose(self):
+        result = p.getLinkState(self.uid, self.ee_link_id, computeForwardKinematics=1)
+        _, _, _, _, frame_pos, frame_quat = result
+        T_world_ee = Transform(Rotation.from_quat(frame_quat), np.array(frame_pos))
+        return self.T_world_base.inverse() * T_world_ee
+
+    def jacobian(self, q):
+        q = np.r_[q, 0.0, 0.0].tolist()
+        q_dot, q_ddot = np.zeros(9).tolist(), np.zeros(9).tolist()
+        linear, _ = p.calculateJacobian(self.uid, 7, [0.0, 0.0, 0.0], q, q_dot, q_ddot)
+        return np.asarray(linear)[:, :7]
+
+
+class PandaGripper(object):
+    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(object):
+    pass
+
+
+class SimPandaEnv(object):
+    def __init__(self, gui, dt=1.0 / 240.0, publish_state=True):
+        self.gui = gui
+        self.dt = dt
+        p.connect(p.GUI if gui else p.DIRECT)
+        p.setAdditionalSearchPath(pybullet_data.getDataPath())
+        p.setTimeStep(dt)
+        p.setGravity(0.0, 0.0, -9.8)
+
+        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.camera = None
+
+        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.dt)
+        for _ in range(steps):
+            self.step()
+
+    def step(self):
+        p.stepSimulation()
+        time.sleep(self.dt)
+
+    def _publish_state(self, event):
+        positions, velocities = self.arm.get_state()
+        msg = JointState()
+        msg.header.stamp = rospy.Time.now()
+        msg.name = ["panda_joint{}".format(i) for i in range(1, 8)]
+        msg.position = positions
+        msg.velocity = velocities
+        self.state_pub.publish(msg)