1import unittest
2import pybullet
3import time
4
5from utils import allclose, dot
6
7
8class TestPybulletMethods(unittest.TestCase):
9
10  def test_import(self):
11    import pybullet as p
12    self.assertGreater(p.getAPIVersion(), 201700000)
13
14  def test_connect_direct(self):
15    import pybullet as p
16    cid = p.connect(p.DIRECT)
17    self.assertEqual(cid, 0)
18    p.disconnect()
19
20  def test_loadurdf(self):
21    import pybullet as p
22    p.connect(p.DIRECT)
23    ob = p.loadURDF("r2d2.urdf")
24    self.assertEqual(ob, 0)
25    p.disconnect()
26
27  def test_rolling_friction(self):
28    import pybullet as p
29    p.connect(p.DIRECT)
30    p.loadURDF("plane.urdf")
31    sphere = p.loadURDF("sphere2.urdf", [0, 0, 1])
32    p.resetBaseVelocity(sphere, linearVelocity=[1, 0, 0])
33    p.changeDynamics(sphere, -1, linearDamping=0, angularDamping=0)
34    #p.changeDynamics(sphere,-1,rollingFriction=0)
35    p.setGravity(0, 0, -10)
36    for i in range(1000):
37      p.stepSimulation()
38    vel = p.getBaseVelocity(sphere)
39    self.assertLess(vel[0][0], 1e-10)
40    self.assertLess(vel[0][1], 1e-10)
41    self.assertLess(vel[0][2], 1e-10)
42    self.assertLess(vel[1][0], 1e-10)
43    self.assertLess(vel[1][1], 1e-10)
44    self.assertLess(vel[1][2], 1e-10)
45    p.disconnect()
46
47
48class TestPybulletJacobian(unittest.TestCase):
49
50  def getMotorJointStates(self, robot):
51    import pybullet as p
52    joint_states = p.getJointStates(robot, range(p.getNumJoints(robot)))
53    joint_infos = [p.getJointInfo(robot, i) for i in range(p.getNumJoints(robot))]
54    joint_states = [j for j, i in zip(joint_states, joint_infos) if i[3] > -1]
55    joint_positions = [state[0] for state in joint_states]
56    joint_velocities = [state[1] for state in joint_states]
57    joint_torques = [state[3] for state in joint_states]
58    return joint_positions, joint_velocities, joint_torques
59
60  def setJointPosition(self, robot, position, kp=1.0, kv=0.3):
61    import pybullet as p
62    num_joints = p.getNumJoints(robot)
63    zero_vec = [0.0] * num_joints
64    if len(position) == num_joints:
65      p.setJointMotorControlArray(robot,
66                                  range(num_joints),
67                                  p.POSITION_CONTROL,
68                                  targetPositions=position,
69                                  targetVelocities=zero_vec,
70                                  positionGains=[kp] * num_joints,
71                                  velocityGains=[kv] * num_joints)
72
73  def testJacobian(self):
74    import pybullet as p
75
76    clid = p.connect(p.SHARED_MEMORY)
77    if (clid < 0):
78      p.connect(p.DIRECT)
79
80    time_step = 0.001
81    gravity_constant = -9.81
82
83    urdfs = [
84        "TwoJointRobot_w_fixedJoints.urdf", "TwoJointRobot_w_fixedJoints.urdf",
85        "kuka_iiwa/model.urdf", "kuka_lwr/kuka.urdf"
86    ]
87    for urdf in urdfs:
88      p.resetSimulation()
89      p.setTimeStep(time_step)
90      p.setGravity(0.0, 0.0, gravity_constant)
91
92      robotId = p.loadURDF(urdf, useFixedBase=True)
93      p.resetBasePositionAndOrientation(robotId, [0, 0, 0], [0, 0, 0, 1])
94      numJoints = p.getNumJoints(robotId)
95      endEffectorIndex = numJoints - 1
96
97      # Set a joint target for the position control and step the sim.
98      self.setJointPosition(robotId, [0.1 * (i % 3) for i in range(numJoints)])
99      p.stepSimulation()
100
101      # Get the joint and link state directly from Bullet.
102      mpos, mvel, mtorq = self.getMotorJointStates(robotId)
103
104      result = p.getLinkState(robotId,
105                              endEffectorIndex,
106                              computeLinkVelocity=1,
107                              computeForwardKinematics=1)
108      link_trn, link_rot, com_trn, com_rot, frame_pos, frame_rot, link_vt, link_vr = result
109      # Get the Jacobians for the CoM of the end-effector link.
110      # Note that in this example com_rot = identity, and we would need to use com_rot.T * com_trn.
111      # The localPosition is always defined in terms of the link frame coordinates.
112
113      zero_vec = [0.0] * len(mpos)
114      jac_t, jac_r = p.calculateJacobian(robotId, endEffectorIndex, com_trn, mpos, zero_vec,
115                                         zero_vec)
116
117      assert (allclose(dot(jac_t, mvel), link_vt))
118      assert (allclose(dot(jac_r, mvel), link_vr))
119    p.disconnect()
120
121
122if __name__ == '__main__':
123  unittest.main()
124