Lecture 2 code examples

Author: Zackory Erickson

.gitignore

@@ -8,6 +8,7 @@ __pycache__/
 *.out
 trained_models*
 *.json
+*.swp
 
 # C extensions
 *.so

README.md

@@ -19,7 +19,7 @@ python3 -m pip install mengine@git+https://github.com/Zackory/mengine.git
 ## Example
 ```bash
 cd examples
-python3 example_push.py
+python3 lecture01_pushing.py
 ```
 
 ![Manipulation Engine](images/mengine.jpg "Manipulation Engine")

examples/lecture02_robots.py

@@ -0,0 +1,24 @@
+import os
+import numpy as np
+import mengine as m
+
+# Create environment and ground plane
+env = m.Env()
+ground = m.Ground()
+
+env.set_gui_camera(distance=2.5)
+
+# Create robots
+robot = m.Robot.Baxter(position=[-2, 0, 0.92], orientation=[0, 0, -np.pi/2])
+
+robot2 = m.Robot.PR2(position=[-0.5, 0, 0], orientation=[0, 0, -np.pi/2])
+
+robot3 = m.Robot.Stretch(position=[0.5, 0, 0], orientation=[0, 0, -np.pi/2])
+
+robot4 = m.Robot.Panda(position=[1.5, 0, 0])
+
+robot5 = m.Robot.Jaco(position=[2, 0, 0])
+
+m.step_simulation(steps=1000, realtime=True)
+
+

examples/lecture02_sensors.py

@@ -0,0 +1,78 @@
+import os
+import numpy as np
+import mengine as m
+import matplotlib.pyplot as plt
+np.set_printoptions(precision=3, suppress=True)
+
+# Create environment and ground plane
+env = m.Env()
+ground = m.Ground()
+
+env.set_gui_camera(distance=2, yaw=30)
+
+# Create table
+table = m.URDF(filename=os.path.join(m.directory, 'table', 'table.urdf'), static=True, position=[0, 0, 0], orientation=[0, 0, 0, 1])
+
+# Create robots
+panda = m.Robot.Panda(position=[0.5, 0, 0.75])
+panda.enable_force_torque_sensor(8) # Create a F/T sensor at the wrist
+# panda.print_joint_info()
+
+stretch = m.Robot.Stretch(position=[0.25, -1, 0], orientation=[0, 0, np.pi/2])
+# stretch.print_joint_info()
+# Rotate the camera
+stretch.set_joint_angles(angles=[-np.pi/8], joints=[21])
+# stretch.set_joint_angles(angles=[-np.pi/4], joints=[20])
+
+# Create a cube to make contact with
+cube = m.Shape(m.Box(half_extents=[0.1]*3), static=False, mass=1, position=[-0.25, 0, 0.85], orientation=[0, 0, 0, 1], rgba=[0, 1, 0, 1])
+
+for i in range(50):
+    # Move the end effector into a cube and print sensor values
+    position, orientation = panda.get_link_pos_orient(panda.end_effector)
+    target_joint_angles = panda.ik(panda.end_effector, target_pos=position+[-0.08, 0, 0], target_orient=orientation, use_current_joint_angles=True)
+    panda.control(target_joint_angles)
+
+    m.step_simulation(realtime=True)
+
+    if i % 20 == 0:
+        # Print joint position, velocities, and torques
+        motor_indices, motor_positions, motor_velocities, motor_torques = panda.get_motor_joint_states()
+        print('Joint angles:', np.array(motor_positions))
+        print('Joint velocities:', np.array(motor_velocities))
+        print('Joint torques:', np.array(motor_torques))
+        print('Force-torque:', panda.get_force_torque_sensor(8))
+        linkA, linkB, posA, posB, contact_distance = panda.get_closest_points(bodyB=cube, linkA=8)
+        print('Proximity (shortest distance from end effector to cube):', contact_distance[0])
+        print('')
+
+# Get the position and orientation for the camera_color_optical_joint
+head_pos, head_orient = stretch.get_link_pos_orient(30)
+# Move link slightly out of the mesh
+head_pos, head_orient = m.multiply_transforms(head_pos, head_orient, [0.01, 0, 0], [0, 0, 0])
+# m.visualize_coordinate_frame(head_pos, head_orient)
+
+# Create a virtual camera for the Stretch's RealSense
+camera = m.Camera(fov=60, camera_width=1920//4, camera_height=1080//4)
+rpy = m.get_euler(head_orient)
+# Align axes between head and camera
+rpy = [rpy[0]-np.pi/2, rpy[1], rpy[2]+np.pi/2]
+camera.set_camera_rpy(look_at_pos=head_pos, distance=-0.001, rpy=np.degrees(rpy))
+
+img, depth, segmentation_mask = camera.get_rgba_depth()
+
+plt.figure(); plt.title('RGB'); plt.imshow(img)
+plt.figure(); plt.title('Depth'); plt.imshow(depth)
+plt.figure(); plt.title('Segmentation Mask'); plt.imshow(segmentation_mask)
+plt.show()
+
+# Capture a point cloud from the camera
+pc, rgba = camera.get_point_cloud(body=panda)
+# pc, rgba = camera.get_point_cloud()
+# Visualize the point cloud
+m.DebugPoints(pc, points_rgb=rgba[:, :3], size=10)
+# Hide the panda to see the visualized point cloud
+panda.set_base_pos_orient([0.5, 0, 100.75])
+
+m.step_simulation(steps=1000, realtime=True)
+

mengine/bodies/baxter.py

@@ -14,12 +14,12 @@ def __init__(self, env, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllab
         # right_gripper_indices = [27, 29] # Gripper actuated joints
         # left_gripper_indices = [49, 51] # Gripper actuated joints
 
-        body = p.loadURDF(os.path.join(directory, 'baxter', 'baxter_custom.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, physicsClientId=env.id)
+        body = p.loadURDF(os.path.join(env.directory, 'baxter', 'baxter_custom.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, physicsClientId=env.id)
         super().__init__(body, env, controllable_joints, end_effector, gripper_joints)
 
         # Recolor robot
         for i in [20, 21, 23, 31, 32, 42, 43, 45, 53, 54]:
-            p.changeVisualShape(self.body, i, rgbaColor=[1.0, 1.0, 1.0, 0.0], physicsClientId=id)
+            p.changeVisualShape(self.body, i, rgbaColor=[1.0, 1.0, 1.0, 0.0], physicsClientId=env.id)
         # Close gripper
         # self.set_gripper_position([0]*2, set_instantly=True)
 

mengine/bodies/jaco.py

@@ -7,6 +7,6 @@ def __init__(self, env, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllab
         controllable_joints = [1, 2, 3, 4, 5, 6, 7] if controllable_joints is None else controllable_joints
         end_effector = 8 # Used to get the pose of the end effector
         gripper_joints = [9, 11, 13] # Gripper actuated joints
-        body = p.loadURDF(os.path.join(directory, 'jaco', 'j2s7s300_gym.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, physicsClientId=env.id)
+        body = p.loadURDF(os.path.join(env.directory, 'jaco', 'j2s7s300_gym.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, physicsClientId=env.id)
         super().__init__(body, env, controllable_joints, end_effector, gripper_joints)
 

mengine/bodies/kinova_gen3.py

@@ -7,6 +7,6 @@ def __init__(self, env, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllab
         controllable_joints = [0, 1, 2, 3, 4, 5, 6] if controllable_joints is None else controllable_joints
         end_effector = 7 # Used to get the pose of the end effector
         gripper_joints = [] # Gripper actuated joints
-        body = p.loadURDF(os.path.join(directory, 'kinova_gen3', 'GEN3_URDF_V12.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, physicsClientId=env.id)
+        body = p.loadURDF(os.path.join(env.directory, 'kinova_gen3', 'GEN3_URDF_V12.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, physicsClientId=env.id)
         super().__init__(body, env, controllable_joints, end_effector, gripper_joints)
 

mengine/bodies/pr2.py

@@ -15,18 +15,18 @@ def __init__(self, env, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllab
         # right_gripper_indices = [57, 58, 59, 60] # Gripper actuated joints
         # left_gripper_indices = [79, 80, 81, 82] # Gripper actuated joints
 
-        body = p.loadURDF(os.path.join(directory, 'PR2', 'pr2_no_torso_lift_tall.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, flags=p.URDF_USE_INERTIA_FROM_FILE, physicsClientId=env.id)
+        body = p.loadURDF(os.path.join(env.directory, 'PR2', 'pr2_no_torso_lift_tall.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, flags=p.URDF_USE_INERTIA_FROM_FILE, physicsClientId=env.id)
         super().__init__(body, env, controllable_joints, end_effector, gripper_joints)
 
         # Recolor robot
         for i in [19, 42, 64]:
-            p.changeVisualShape(self.body, i, rgbaColor=[1.0, 1.0, 1.0, 1.0], physicsClientId=id)
+            p.changeVisualShape(self.body, i, rgbaColor=[1.0, 1.0, 1.0, 1.0], physicsClientId=env.id)
         for i in [43, 46, 49, 58, 60, 65, 68, 71, 80, 82]:
-            p.changeVisualShape(self.body, i, rgbaColor=[0.4, 0.4, 0.4, 1.0], physicsClientId=id)
+            p.changeVisualShape(self.body, i, rgbaColor=[0.4, 0.4, 0.4, 1.0], physicsClientId=env.id)
         for i in [45, 51, 67, 73]:
-            p.changeVisualShape(self.body, i, rgbaColor=[0.7, 0.7, 0.7, 1.0], physicsClientId=id)
-        p.changeVisualShape(self.body, 20, rgbaColor=[0.8, 0.8, 0.8, 1.0], physicsClientId=id)
-        p.changeVisualShape(self.body, 40, rgbaColor=[0.6, 0.6, 0.6, 1.0], physicsClientId=id)
+            p.changeVisualShape(self.body, i, rgbaColor=[0.7, 0.7, 0.7, 1.0], physicsClientId=env.id)
+        p.changeVisualShape(self.body, 20, rgbaColor=[0.8, 0.8, 0.8, 1.0], physicsClientId=env.id)
+        p.changeVisualShape(self.body, 40, rgbaColor=[0.6, 0.6, 0.6, 1.0], physicsClientId=env.id)
         # Close gripper
         # self.set_gripper_position([0]*2, set_instantly=True)
 

mengine/bodies/stretch.py

@@ -7,7 +7,7 @@ def __init__(self, env, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllab
         controllable_joints = [3, 5, 6, 7, 8, 9] if controllable_joints is None else controllable_joints
         end_effector = 15 # Used to get the pose of the end effector
         gripper_joints = [11, 13] # Gripper actuated joints
-        body = p.loadURDF(os.path.join(directory, 'stretch', 'stretch_uncalibrated.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, physicsClientId=env.id)
+        body = p.loadURDF(os.path.join(env.directory, 'stretch', 'stretch_uncalibrated.urdf'), useFixedBase=fixed_base, basePosition=position, baseOrientation=orientation, physicsClientId=env.id)
         super().__init__(body, env, controllable_joints, end_effector, gripper_joints)
 
         # Recolor robot
@@ -16,11 +16,11 @@ def __init__(self, env, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllab
         dark_gray = [0.4, 0.4, 0.4, 1]
         black = [0.251, 0.251, 0.251, 1]
         for i in [0, 1, 2, 4, 5, 6, 7, 8, 9, 10, 12, 14, 19, 20, 23, 35]:
-            p.changeVisualShape(self.body, i, rgbaColor=black, physicsClientId=id)
+            p.changeVisualShape(self.body, i, rgbaColor=black, physicsClientId=env.id)
         for i in [-1, 11, 13, 21, 16, 17, 18, 33, 34]:
-            p.changeVisualShape(self.body, i, rgbaColor=gray, physicsClientId=id)
+            p.changeVisualShape(self.body, i, rgbaColor=gray, physicsClientId=env.id)
         # for i in [16, 17, 18, 33, 34]:
         #     p.changeVisualShape(self.body, i, rgbaColor=white, physicsClientId=id)
         for i in [3, 32]:
-            p.changeVisualShape(self.body, i, rgbaColor=dark_gray, physicsClientId=id)
+            p.changeVisualShape(self.body, i, rgbaColor=dark_gray, physicsClientId=env.id)
 

mengine/env.py

@@ -7,6 +7,11 @@
 from .util import Util
 from .bodies.body import Body
 from .bodies.panda import Panda
+from .bodies.jaco import Jaco
+from .bodies.baxter import Baxter
+from .bodies.pr2 import PR2
+from .bodies.stretch import Stretch
+from .bodies.kinova_gen3 import KinovaGen3
 
 envir = None
 directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'assets')
@@ -100,6 +105,31 @@ def __init__(self, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllable_jo
             env = env if env is not None else envir
             super().__init__(env, position, get_quaternion(orientation), controllable_joints, fixed_base)
 
+    class Jaco(Jaco):
+        def __init__(self, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllable_joints=None, fixed_base=True, env=None):
+            env = env if env is not None else envir
+            super().__init__(env, position, get_quaternion(orientation), controllable_joints, fixed_base)
+
+    class Baxter(Baxter):
+        def __init__(self, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllable_joints=None, fixed_base=True, env=None):
+            env = env if env is not None else envir
+            super().__init__(env, position, get_quaternion(orientation), controllable_joints, fixed_base)
+
+    class PR2(PR2):
+        def __init__(self, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllable_joints=None, fixed_base=True, env=None):
+            env = env if env is not None else envir
+            super().__init__(env, position, get_quaternion(orientation), controllable_joints, fixed_base)
+
+    class Stretch(Stretch):
+        def __init__(self, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllable_joints=None, fixed_base=True, env=None):
+            env = env if env is not None else envir
+            super().__init__(env, position, get_quaternion(orientation), controllable_joints, fixed_base)
+
+    class KinovaGen3(KinovaGen3):
+        def __init__(self, position=[0, 0, 0], orientation=[0, 0, 0, 1], controllable_joints=None, fixed_base=True, env=None):
+            env = env if env is not None else envir
+            super().__init__(env, position, get_quaternion(orientation), controllable_joints, fixed_base)
+
 class Camera:
     def __init__(self, camera_pos=[0.5, -0.5, 1.5], look_at_pos=[0, 0, 0.75], fov=60, camera_width=1920//4, camera_height=1080//4, env=None):
         self.env = env if env is not None else envir
@@ -141,13 +171,14 @@ def get_point_cloud(self, body=None):
         pixels = np.stack([x, y, z, h], axis=1)
 
         # Filter point cloud to only include points on the target body
-        pixels = pixels[segmentation_mask == body.body]
-        z = z[segmentation_mask == body.body]
-        rgba = rgba[segmentation_mask == body.body]
+        if body is not None:
+            pixels = pixels[segmentation_mask == body.body]
+            z = z[segmentation_mask == body.body]
+            rgba = rgba[segmentation_mask == body.body]
 
         # filter out "infinite" depths
-        pixels = pixels[z < 0.99]
-        rgba = rgba[z < 0.99]
+        pixels = pixels[z < 20]
+        rgba = rgba[z < 20]
         pixels[:, 2] = 2 * pixels[:, 2] - 1
 
         # turn pixels to world coordinates
@@ -182,7 +213,7 @@ def get_difference_quaternion(q1, q2, env=None):
 def quaternion_product(q1, q2, env=None):
     env = env if env is not None else envir
     # Return Hamilton product of 2 quaternions
-    return p.multiplyTransforms([0, 0, 0], get_quaternion(q1), [0, 0, 0], q2, physicsClientId=env.id)[1]
+    return p.multiplyTransforms([0, 0, 0], get_quaternion(q1), [0, 0, 0], get_quaternion(q2), physicsClientId=env.id)[1]
 
 def multiply_transforms(p1, q1, p2, q2, env=None):
     env = env if env is not None else envir

setup.py

@@ -18,7 +18,7 @@
     packages=find_packages(),
     python_requires='>=3',
     # install_requires=['gym>=0.2.3', 'pybullet @ git+https://github.com/Zackory/bullet3.git@pybullet_3_0_9#egg=pybullet', 'numpy', 'keras==2.3.0', 'tensorflow==1.14.0', 'h5py==2.10.0', 'smplx', 'trimesh', 'ray[rllib]', 'numpngw', 'tensorflow-probability==0.7.0', 'matplotlib'] + ['screeninfo==0.6.1' if sys.version_info >= (3, 6) else 'screeninfo==0.2'],
-    install_requires=['pybullet', 'numpy', 'scipy', 'numpngw'] + ['screeninfo==0.6.1' if sys.version_info >= (3, 6) else 'screeninfo==0.2'],
+    install_requires=['pybullet', 'numpy', 'scipy', 'numpngw', 'matplotlib'] + ['screeninfo==0.6.1' if sys.version_info >= (3, 6) else 'screeninfo==0.2'],
     # description='Physics simulation for assistive robotics and human-robot interaction.',
     # long_description=long_description,
     long_description_content_type="text/markdown",