Jacky Liang commited on
Commit
922e0e8
1 Parent(s): 5046e86

correctly handling state

Browse files
Files changed (2) hide show
  1. app.py +24 -14
  2. sim.py +68 -67
app.py CHANGED
@@ -16,7 +16,7 @@ from md_logger import MarkdownLogger
16
 
17
 
18
  class DemoRunner:
19
-
20
  def __init__(self):
21
  self._cfg = OmegaConf.to_container(OmegaConf.load('cfg.yaml'), resolve=True)
22
  self._env = None
@@ -69,12 +69,6 @@ class DemoRunner:
69
  return lmp_tabletop_ui
70
 
71
  def setup(self, api_key, model_name, n_blocks, n_bowls):
72
- if not api_key:
73
- return 'Please enter your OpenAI API key!', None
74
-
75
- if n_blocks + n_bowls == 0:
76
- return 'Please select at least one object!', None
77
-
78
  openai.api_key = api_key
79
  self._model_name = model_name
80
 
@@ -112,16 +106,32 @@ class DemoRunner:
112
  return self._md_logger.get_log(), self._env.get_camera_image(), video_file_name
113
 
114
 
115
- if __name__ == '__main__':
 
 
 
 
 
 
116
  demo_runner = DemoRunner()
117
- demo = gr.Blocks()
 
 
118
 
 
 
 
 
 
 
119
  with open('README.md', 'r') as f:
120
  for _ in range(12):
121
  next(f)
122
  readme_text = f.read()
123
 
124
- with demo:
 
 
125
  gr.Markdown(readme_text)
126
  gr.Markdown('# Interactive Demo')
127
  with gr.Row():
@@ -148,13 +158,13 @@ if __name__ == '__main__':
148
  video_run = gr.Video(label='Video of Last Instruction')
149
 
150
  btn_setup.click(
151
- demo_runner.setup,
152
  inputs=[inp_api_key, inp_model_name, inp_n_blocks, inp_n_bowls],
153
- outputs=[info_setup, img_setup]
154
  )
155
  btn_run.click(
156
- demo_runner.run,
157
- inputs=[inp_instruction],
158
  outputs=[info_run, img_setup, video_run]
159
  )
160
 
 
16
 
17
 
18
  class DemoRunner:
19
+
20
  def __init__(self):
21
  self._cfg = OmegaConf.to_container(OmegaConf.load('cfg.yaml'), resolve=True)
22
  self._env = None
 
69
  return lmp_tabletop_ui
70
 
71
  def setup(self, api_key, model_name, n_blocks, n_bowls):
 
 
 
 
 
 
72
  openai.api_key = api_key
73
  self._model_name = model_name
74
 
 
106
  return self._md_logger.get_log(), self._env.get_camera_image(), video_file_name
107
 
108
 
109
+ def setup(api_key, model_name, n_blocks, n_bowls):
110
+ if not api_key:
111
+ return 'Please enter your OpenAI API key!', None
112
+
113
+ if n_blocks + n_bowls == 0:
114
+ return 'Please select at least one object!', None
115
+
116
  demo_runner = DemoRunner()
117
+
118
+ info, img = demo_runner.setup(api_key, model_name, n_blocks, n_bowls)
119
+ return info, img, demo_runner
120
 
121
+
122
+ def run(instruction, demo_runner):
123
+ return demo_runner.run(instruction)
124
+
125
+
126
+ if __name__ == '__main__':
127
  with open('README.md', 'r') as f:
128
  for _ in range(12):
129
  next(f)
130
  readme_text = f.read()
131
 
132
+ with gr.Blocks() as demo:
133
+ state = gr.State(None)
134
+
135
  gr.Markdown(readme_text)
136
  gr.Markdown('# Interactive Demo')
137
  with gr.Row():
 
158
  video_run = gr.Video(label='Video of Last Instruction')
159
 
160
  btn_setup.click(
161
+ setup,
162
  inputs=[inp_api_key, inp_model_name, inp_n_blocks, inp_n_bowls],
163
+ outputs=[info_setup, img_setup, state]
164
  )
165
  btn_run.click(
166
+ run,
167
+ inputs=[inp_instruction, state],
168
  outputs=[info_run, img_setup, video_run]
169
  )
170
 
sim.py CHANGED
@@ -1,4 +1,5 @@
1
  import pybullet
 
2
  import pybullet_data
3
  import threading
4
  from time import sleep
@@ -7,26 +8,28 @@ import os
7
  from consts import BOUNDS, COLORS, PIXEL_SIZE, CORNER_POS
8
  from shapely.geometry import box
9
 
 
10
  # Gripper (Robotiq 2F85) code
11
  class Robotiq2F85:
12
  """Gripper handling for Robotiq 2F85."""
13
 
14
- def __init__(self, robot, tool):
15
  self.robot = robot
16
  self.tool = tool
 
17
  pos = [0.1339999999999999, -0.49199999999872496, 0.5]
18
- rot = pybullet.getQuaternionFromEuler([np.pi, 0, np.pi])
19
  urdf = 'robotiq_2f_85/robotiq_2f_85.urdf'
20
- self.body = pybullet.loadURDF(urdf, pos, rot)
21
- self.n_joints = pybullet.getNumJoints(self.body)
22
  self.activated = False
23
 
24
  # Connect gripper base to robot tool.
25
- pybullet.createConstraint(self.robot, tool, self.body, 0, jointType=pybullet.JOINT_FIXED, jointAxis=[0, 0, 0], parentFramePosition=[0, 0, 0], childFramePosition=[0, 0, -0.07], childFrameOrientation=pybullet.getQuaternionFromEuler([0, 0, np.pi / 2]))
26
 
27
  # Set friction coefficients for gripper fingers.
28
- for i in range(pybullet.getNumJoints(self.body)):
29
- pybullet.changeDynamics(self.body, i, lateralFriction=10.0, spinningFriction=1.0, rollingFriction=1.0, frictionAnchor=True)
30
 
31
  # Start thread to handle additional gripper constraints.
32
  self.motor_joint = 1
@@ -39,22 +42,22 @@ class Robotiq2F85:
39
  def step(self):
40
  while True:
41
  try:
42
- currj = [pybullet.getJointState(self.body, i)[0] for i in range(self.n_joints)]
43
  indj = [6, 3, 8, 5, 10]
44
  targj = [currj[1], -currj[1], -currj[1], currj[1], currj[1]]
45
- pybullet.setJointMotorControlArray(self.body, indj, pybullet.POSITION_CONTROL, targj, positionGains=np.ones(5))
46
  except:
47
  return
48
  sleep(0.001)
49
 
50
  # Close gripper fingers.
51
  def activate(self):
52
- pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=1, force=10)
53
  self.activated = True
54
 
55
  # Open gripper fingers.
56
  def release(self):
57
- pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=-1, force=10)
58
  self.activated = False
59
 
60
  # If activated and object in gripper: check object contact.
@@ -75,7 +78,7 @@ class Robotiq2F85:
75
  def external_contact(self, body=None):
76
  if body is None:
77
  body = self.body
78
- pts = pybullet.getContactPoints(bodyA=body)
79
  pts = [pt for pt in pts if pt[2] != self.body]
80
  return len(pts) > 0 # pylint: disable=g-explicit-length-test
81
 
@@ -86,17 +89,17 @@ class Robotiq2F85:
86
  return success
87
 
88
  def grasp_width(self):
89
- lpad = np.array(pybullet.getLinkState(self.body, 4)[0])
90
- rpad = np.array(pybullet.getLinkState(self.body, 9)[0])
91
  dist = np.linalg.norm(lpad - rpad) - 0.047813
92
  return dist
93
 
94
  def check_proximity(self):
95
- ee_pos = np.array(pybullet.getLinkState(self.robot, self.tool)[0])
96
- tool_pos = np.array(pybullet.getLinkState(self.body, 0)[0])
97
  vec = (tool_pos - ee_pos) / np.linalg.norm((tool_pos - ee_pos))
98
  ee_targ = ee_pos + vec
99
- ray_data = pybullet.rayTest(ee_pos, ee_targ)[0]
100
  obj, link, ray_frac = ray_data[0], ray_data[1], ray_data[2]
101
  return obj, link, ray_frac
102
 
@@ -108,16 +111,15 @@ class PickPlaceEnv():
108
  self.dt = 1/480
109
  self.sim_step = 0
110
 
111
- # Configure and start PyBullet.
112
- # python3 -m pybullet_utils.runServer
113
- # pybullet.connect(pybullet.SHARED_MEMORY) # pybullet.GUI for local GUI.
114
- pybullet.connect(pybullet.DIRECT) # pybullet.GUI for local GUI.
115
- pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_GUI, 0)
116
- pybullet.setPhysicsEngineParameter(enableFileCaching=0)
117
  assets_path = os.path.dirname(os.path.abspath(""))
118
- pybullet.setAdditionalSearchPath(assets_path)
119
- pybullet.setAdditionalSearchPath(pybullet_data.getDataPath())
120
- pybullet.setTimeStep(self.dt)
121
 
122
  self.home_joints = (np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 3 * np.pi / 2, 0) # Joint angles: (J0, J1, J2, J3, J4, J5).
123
  self.home_ee_euler = (np.pi, 0, np.pi) # (RX, RY, RZ) rotation in Euler angles.
@@ -130,36 +132,36 @@ class PickPlaceEnv():
130
  self.high_frame_rate = high_frame_rate
131
 
132
  def reset(self, object_list):
133
- pybullet.resetSimulation(pybullet.RESET_USE_DEFORMABLE_WORLD)
134
- pybullet.setGravity(0, 0, -9.8)
135
  self.cache_video = []
136
 
137
  # Temporarily disable rendering to load URDFs faster.
138
- pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 0)
139
 
140
  # Add robot.
141
- pybullet.loadURDF("plane.urdf", [0, 0, -0.001])
142
- self.robot_id = pybullet.loadURDF("ur5e/ur5e.urdf", [0, 0, 0], flags=pybullet.URDF_USE_MATERIAL_COLORS_FROM_MTL)
143
- self.ghost_id = pybullet.loadURDF("ur5e/ur5e.urdf", [0, 0, -10]) # For forward kinematics.
144
- self.joint_ids = [pybullet.getJointInfo(self.robot_id, i) for i in range(pybullet.getNumJoints(self.robot_id))]
145
- self.joint_ids = [j[0] for j in self.joint_ids if j[2] == pybullet.JOINT_REVOLUTE]
146
 
147
  # Move robot to home configuration.
148
  for i in range(len(self.joint_ids)):
149
- pybullet.resetJointState(self.robot_id, self.joint_ids[i], self.home_joints[i])
150
 
151
  # Add gripper.
152
  if self.gripper is not None:
153
  while self.gripper.constraints_thread.is_alive():
154
  self.constraints_thread_active = False
155
- self.gripper = Robotiq2F85(self.robot_id, self.ee_link_id)
156
  self.gripper.release()
157
 
158
  # Add workspace.
159
- plane_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])
160
- plane_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])
161
- plane_id = pybullet.createMultiBody(0, plane_shape, plane_visual, basePosition=[0, -0.5, 0])
162
- pybullet.changeVisualShape(plane_id, -1, rgbaColor=[0.2, 0.2, 0.2, 1.0])
163
 
164
  # Load objects according to config.
165
  self.object_list = object_list
@@ -186,21 +188,20 @@ class PickPlaceEnv():
186
  object_type = obj_name.split(' ')[1]
187
  object_position = rand_xyz.squeeze()
188
  if object_type == 'block':
189
- object_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])
190
- object_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])
191
- object_id = pybullet.createMultiBody(0.01, object_shape, object_visual, basePosition=object_position)
192
  elif object_type == 'bowl':
193
  object_position[2] = 0
194
- object_id = pybullet.loadURDF("bowl/bowl.urdf", object_position, useFixedBase=1)
195
- pybullet.changeVisualShape(object_id, -1, rgbaColor=object_color)
196
  self.obj_name_to_id[obj_name] = object_id
197
-
198
 
199
  # Re-enable rendering.
200
- pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 1)
201
 
202
  for _ in range(200):
203
- pybullet.stepSimulation()
204
 
205
  # record object positions at reset
206
  self.init_pos = {name: self.get_obj_pos(name) for name in object_list}
@@ -209,25 +210,25 @@ class PickPlaceEnv():
209
 
210
  def servoj(self, joints):
211
  """Move to target joint positions with position control."""
212
- pybullet.setJointMotorControlArray(
213
  bodyIndex=self.robot_id,
214
  jointIndices=self.joint_ids,
215
- controlMode=pybullet.POSITION_CONTROL,
216
  targetPositions=joints,
217
  positionGains=[0.01]*6)
218
 
219
  def movep(self, position):
220
  """Move to target end effector position."""
221
- joints = pybullet.calculateInverseKinematics(
222
  bodyUniqueId=self.robot_id,
223
  endEffectorLinkIndex=self.tip_link_id,
224
  targetPosition=position,
225
- targetOrientation=pybullet.getQuaternionFromEuler(self.home_ee_euler),
226
  maxNumIterations=100)
227
  self.servoj(joints)
228
 
229
  def get_ee_pos(self):
230
- ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])
231
  return ee_xyz
232
 
233
  def step(self, action=None):
@@ -306,17 +307,17 @@ class PickPlaceEnv():
306
 
307
  def set_alpha_transparency(self, alpha: float) -> None:
308
  for id in range(20):
309
- visual_shape_data = pybullet.getVisualShapeData(id)
310
  for i in range(len(visual_shape_data)):
311
  object_id, link_index, _, _, _, _, _, rgba_color = visual_shape_data[i]
312
  rgba_color = list(rgba_color[0:3]) + [alpha]
313
- pybullet.changeVisualShape(
314
  self.robot_id, linkIndex=i, rgbaColor=rgba_color)
315
- pybullet.changeVisualShape(
316
  self.gripper.body, linkIndex=i, rgbaColor=rgba_color)
317
 
318
  def step_sim_and_render(self):
319
- pybullet.stepSimulation()
320
  self.sim_step += 1
321
 
322
  interval = 40 if self.high_frame_rate else 60
@@ -346,7 +347,7 @@ class PickPlaceEnv():
346
  # Get heightmaps and colormaps.
347
  points = self.get_pointcloud(depth, intrinsics)
348
  position = np.float32(position).reshape(3, 1)
349
- rotation = pybullet.getMatrixFromQuaternion(orientation)
350
  rotation = np.float32(rotation).reshape(3, 3)
351
  transform = np.eye(4)
352
  transform[:3, :] = np.hstack((rotation, position))
@@ -363,37 +364,37 @@ class PickPlaceEnv():
363
  # Camera parameters.
364
  position = (0, -0.85, 0.4)
365
  orientation = (np.pi / 4 + np.pi / 48, np.pi, np.pi)
366
- orientation = pybullet.getQuaternionFromEuler(orientation)
367
  zrange = (0.01, 10.)
368
  noise=True
369
 
370
  # OpenGL camera settings.
371
  lookdir = np.float32([0, 0, 1]).reshape(3, 1)
372
  updir = np.float32([0, -1, 0]).reshape(3, 1)
373
- rotation = pybullet.getMatrixFromQuaternion(orientation)
374
  rotm = np.float32(rotation).reshape(3, 3)
375
  lookdir = (rotm @ lookdir).reshape(-1)
376
  updir = (rotm @ updir).reshape(-1)
377
  lookat = position + lookdir
378
  focal_len = intrinsics[0]
379
  znear, zfar = (0.01, 10.)
380
- viewm = pybullet.computeViewMatrix(position, lookat, updir)
381
  fovh = (image_size[0] / 2) / focal_len
382
  fovh = 180 * np.arctan(fovh) * 2 / np.pi
383
 
384
  # Notes: 1) FOV is vertical FOV 2) aspect must be float
385
  aspect_ratio = image_size[1] / image_size[0]
386
- projm = pybullet.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)
387
 
388
  # Render with OpenGL camera settings.
389
- _, _, color, depth, segm = pybullet.getCameraImage(
390
  width=image_size[1],
391
  height=image_size[0],
392
  viewMatrix=viewm,
393
  projectionMatrix=projm,
394
  shadow=1,
395
- flags=pybullet.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,
396
- renderer=pybullet.ER_BULLET_HARDWARE_OPENGL)
397
 
398
  # Get color image.
399
  color_image_size = (image_size[0], image_size[1], 4)
@@ -533,13 +534,13 @@ class PickPlaceEnv():
533
  position = np.float32(np.array(CORNER_POS[obj_name]))
534
  else:
535
  pick_id = self.get_obj_id(obj_name)
536
- pose = pybullet.getBasePositionAndOrientation(pick_id)
537
  position = np.float32(pose[0])
538
  return position
539
 
540
  def get_bounding_box(self, obj_name):
541
  obj_id = self.get_obj_id(obj_name)
542
- return pybullet.getAABB(obj_id)
543
 
544
 
545
  class LMP_wrapper():
 
1
  import pybullet
2
+ from pybullet_utils.bullet_client import BulletClient
3
  import pybullet_data
4
  import threading
5
  from time import sleep
 
8
  from consts import BOUNDS, COLORS, PIXEL_SIZE, CORNER_POS
9
  from shapely.geometry import box
10
 
11
+
12
  # Gripper (Robotiq 2F85) code
13
  class Robotiq2F85:
14
  """Gripper handling for Robotiq 2F85."""
15
 
16
+ def __init__(self, robot, tool, p):
17
  self.robot = robot
18
  self.tool = tool
19
+ self._p = p
20
  pos = [0.1339999999999999, -0.49199999999872496, 0.5]
21
+ rot = self._p.getQuaternionFromEuler([np.pi, 0, np.pi])
22
  urdf = 'robotiq_2f_85/robotiq_2f_85.urdf'
23
+ self.body = self._p.loadURDF(urdf, pos, rot)
24
+ self.n_joints = self._p.getNumJoints(self.body)
25
  self.activated = False
26
 
27
  # Connect gripper base to robot tool.
28
+ self._p.createConstraint(self.robot, tool, self.body, 0, jointType=self._p.JOINT_FIXED, jointAxis=[0, 0, 0], parentFramePosition=[0, 0, 0], childFramePosition=[0, 0, -0.07], childFrameOrientation=self._p.getQuaternionFromEuler([0, 0, np.pi / 2]))
29
 
30
  # Set friction coefficients for gripper fingers.
31
+ for i in range(self._p.getNumJoints(self.body)):
32
+ self._p.changeDynamics(self.body, i, lateralFriction=10.0, spinningFriction=1.0, rollingFriction=1.0, frictionAnchor=True)
33
 
34
  # Start thread to handle additional gripper constraints.
35
  self.motor_joint = 1
 
42
  def step(self):
43
  while True:
44
  try:
45
+ currj = [self._p.getJointState(self.body, i)[0] for i in range(self.n_joints)]
46
  indj = [6, 3, 8, 5, 10]
47
  targj = [currj[1], -currj[1], -currj[1], currj[1], currj[1]]
48
+ self._p.setJointMotorControlArray(self.body, indj, self._p.POSITION_CONTROL, targj, positionGains=np.ones(5))
49
  except:
50
  return
51
  sleep(0.001)
52
 
53
  # Close gripper fingers.
54
  def activate(self):
55
+ self._p.setJointMotorControl2(self.body, self.motor_joint, self._p.VELOCITY_CONTROL, targetVelocity=1, force=10)
56
  self.activated = True
57
 
58
  # Open gripper fingers.
59
  def release(self):
60
+ self._p.setJointMotorControl2(self.body, self.motor_joint, self._p.VELOCITY_CONTROL, targetVelocity=-1, force=10)
61
  self.activated = False
62
 
63
  # If activated and object in gripper: check object contact.
 
78
  def external_contact(self, body=None):
79
  if body is None:
80
  body = self.body
81
+ pts = self._p.getContactPoints(bodyA=body)
82
  pts = [pt for pt in pts if pt[2] != self.body]
83
  return len(pts) > 0 # pylint: disable=g-explicit-length-test
84
 
 
89
  return success
90
 
91
  def grasp_width(self):
92
+ lpad = np.array(self._p.getLinkState(self.body, 4)[0])
93
+ rpad = np.array(self._p.getLinkState(self.body, 9)[0])
94
  dist = np.linalg.norm(lpad - rpad) - 0.047813
95
  return dist
96
 
97
  def check_proximity(self):
98
+ ee_pos = np.array(self._p.getLinkState(self.robot, self.tool)[0])
99
+ tool_pos = np.array(self._p.getLinkState(self.body, 0)[0])
100
  vec = (tool_pos - ee_pos) / np.linalg.norm((tool_pos - ee_pos))
101
  ee_targ = ee_pos + vec
102
+ ray_data = self._p.rayTest(ee_pos, ee_targ)[0]
103
  obj, link, ray_frac = ray_data[0], ray_data[1], ray_data[2]
104
  return obj, link, ray_frac
105
 
 
111
  self.dt = 1/480
112
  self.sim_step = 0
113
 
114
+ # Configure and start PyBullet
115
+ # self._p = pybullet.connect(pybullet.DIRECT)
116
+ self._p = BulletClient(connection_mode=pybullet.DIRECT)
117
+ self._p.configureDebugVisualizer(self._p.COV_ENABLE_GUI, 0)
118
+ self._p.setPhysicsEngineParameter(enableFileCaching=0)
 
119
  assets_path = os.path.dirname(os.path.abspath(""))
120
+ self._p.setAdditionalSearchPath(assets_path)
121
+ self._p.setAdditionalSearchPath(pybullet_data.getDataPath())
122
+ self._p.setTimeStep(self.dt)
123
 
124
  self.home_joints = (np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 3 * np.pi / 2, 0) # Joint angles: (J0, J1, J2, J3, J4, J5).
125
  self.home_ee_euler = (np.pi, 0, np.pi) # (RX, RY, RZ) rotation in Euler angles.
 
132
  self.high_frame_rate = high_frame_rate
133
 
134
  def reset(self, object_list):
135
+ self._p.resetSimulation(self._p.RESET_USE_DEFORMABLE_WORLD)
136
+ self._p.setGravity(0, 0, -9.8)
137
  self.cache_video = []
138
 
139
  # Temporarily disable rendering to load URDFs faster.
140
+ self._p.configureDebugVisualizer(self._p.COV_ENABLE_RENDERING, 0)
141
 
142
  # Add robot.
143
+ self._p.loadURDF("plane.urdf", [0, 0, -0.001])
144
+ self.robot_id = self._p.loadURDF("ur5e/ur5e.urdf", [0, 0, 0], flags=self._p.URDF_USE_MATERIAL_COLORS_FROM_MTL)
145
+ self.ghost_id = self._p.loadURDF("ur5e/ur5e.urdf", [0, 0, -10]) # For forward kinematics.
146
+ self.joint_ids = [self._p.getJointInfo(self.robot_id, i) for i in range(self._p.getNumJoints(self.robot_id))]
147
+ self.joint_ids = [j[0] for j in self.joint_ids if j[2] == self._p.JOINT_REVOLUTE]
148
 
149
  # Move robot to home configuration.
150
  for i in range(len(self.joint_ids)):
151
+ self._p.resetJointState(self.robot_id, self.joint_ids[i], self.home_joints[i])
152
 
153
  # Add gripper.
154
  if self.gripper is not None:
155
  while self.gripper.constraints_thread.is_alive():
156
  self.constraints_thread_active = False
157
+ self.gripper = Robotiq2F85(self.robot_id, self.ee_link_id, self._p)
158
  self.gripper.release()
159
 
160
  # Add workspace.
161
+ plane_shape = self._p.createCollisionShape(self._p.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])
162
+ plane_visual = self._p.createVisualShape(self._p.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])
163
+ plane_id = self._p.createMultiBody(0, plane_shape, plane_visual, basePosition=[0, -0.5, 0])
164
+ self._p.changeVisualShape(plane_id, -1, rgbaColor=[0.2, 0.2, 0.2, 1.0])
165
 
166
  # Load objects according to config.
167
  self.object_list = object_list
 
188
  object_type = obj_name.split(' ')[1]
189
  object_position = rand_xyz.squeeze()
190
  if object_type == 'block':
191
+ object_shape = self._p.createCollisionShape(self._p.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])
192
+ object_visual = self._p.createVisualShape(self._p.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])
193
+ object_id = self._p.createMultiBody(0.01, object_shape, object_visual, basePosition=object_position)
194
  elif object_type == 'bowl':
195
  object_position[2] = 0
196
+ object_id = self._p.loadURDF("bowl/bowl.urdf", object_position, useFixedBase=1)
197
+ self._p.changeVisualShape(object_id, -1, rgbaColor=object_color)
198
  self.obj_name_to_id[obj_name] = object_id
 
199
 
200
  # Re-enable rendering.
201
+ self._p.configureDebugVisualizer(self._p.COV_ENABLE_RENDERING, 1)
202
 
203
  for _ in range(200):
204
+ self._p.stepSimulation()
205
 
206
  # record object positions at reset
207
  self.init_pos = {name: self.get_obj_pos(name) for name in object_list}
 
210
 
211
  def servoj(self, joints):
212
  """Move to target joint positions with position control."""
213
+ self._p.setJointMotorControlArray(
214
  bodyIndex=self.robot_id,
215
  jointIndices=self.joint_ids,
216
+ controlMode=self._p.POSITION_CONTROL,
217
  targetPositions=joints,
218
  positionGains=[0.01]*6)
219
 
220
  def movep(self, position):
221
  """Move to target end effector position."""
222
+ joints = self._p.calculateInverseKinematics(
223
  bodyUniqueId=self.robot_id,
224
  endEffectorLinkIndex=self.tip_link_id,
225
  targetPosition=position,
226
+ targetOrientation=self._p.getQuaternionFromEuler(self.home_ee_euler),
227
  maxNumIterations=100)
228
  self.servoj(joints)
229
 
230
  def get_ee_pos(self):
231
+ ee_xyz = np.float32(self._p.getLinkState(self.robot_id, self.tip_link_id)[0])
232
  return ee_xyz
233
 
234
  def step(self, action=None):
 
307
 
308
  def set_alpha_transparency(self, alpha: float) -> None:
309
  for id in range(20):
310
+ visual_shape_data = self._p.getVisualShapeData(id)
311
  for i in range(len(visual_shape_data)):
312
  object_id, link_index, _, _, _, _, _, rgba_color = visual_shape_data[i]
313
  rgba_color = list(rgba_color[0:3]) + [alpha]
314
+ self._p.changeVisualShape(
315
  self.robot_id, linkIndex=i, rgbaColor=rgba_color)
316
+ self._p.changeVisualShape(
317
  self.gripper.body, linkIndex=i, rgbaColor=rgba_color)
318
 
319
  def step_sim_and_render(self):
320
+ self._p.stepSimulation()
321
  self.sim_step += 1
322
 
323
  interval = 40 if self.high_frame_rate else 60
 
347
  # Get heightmaps and colormaps.
348
  points = self.get_pointcloud(depth, intrinsics)
349
  position = np.float32(position).reshape(3, 1)
350
+ rotation = self._p.getMatrixFromQuaternion(orientation)
351
  rotation = np.float32(rotation).reshape(3, 3)
352
  transform = np.eye(4)
353
  transform[:3, :] = np.hstack((rotation, position))
 
364
  # Camera parameters.
365
  position = (0, -0.85, 0.4)
366
  orientation = (np.pi / 4 + np.pi / 48, np.pi, np.pi)
367
+ orientation = self._p.getQuaternionFromEuler(orientation)
368
  zrange = (0.01, 10.)
369
  noise=True
370
 
371
  # OpenGL camera settings.
372
  lookdir = np.float32([0, 0, 1]).reshape(3, 1)
373
  updir = np.float32([0, -1, 0]).reshape(3, 1)
374
+ rotation = self._p.getMatrixFromQuaternion(orientation)
375
  rotm = np.float32(rotation).reshape(3, 3)
376
  lookdir = (rotm @ lookdir).reshape(-1)
377
  updir = (rotm @ updir).reshape(-1)
378
  lookat = position + lookdir
379
  focal_len = intrinsics[0]
380
  znear, zfar = (0.01, 10.)
381
+ viewm = self._p.computeViewMatrix(position, lookat, updir)
382
  fovh = (image_size[0] / 2) / focal_len
383
  fovh = 180 * np.arctan(fovh) * 2 / np.pi
384
 
385
  # Notes: 1) FOV is vertical FOV 2) aspect must be float
386
  aspect_ratio = image_size[1] / image_size[0]
387
+ projm = self._p.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)
388
 
389
  # Render with OpenGL camera settings.
390
+ _, _, color, depth, segm = self._p.getCameraImage(
391
  width=image_size[1],
392
  height=image_size[0],
393
  viewMatrix=viewm,
394
  projectionMatrix=projm,
395
  shadow=1,
396
+ flags=self._p.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,
397
+ renderer=self._p.ER_BULLET_HARDWARE_OPENGL)
398
 
399
  # Get color image.
400
  color_image_size = (image_size[0], image_size[1], 4)
 
534
  position = np.float32(np.array(CORNER_POS[obj_name]))
535
  else:
536
  pick_id = self.get_obj_id(obj_name)
537
+ pose = self._p.getBasePositionAndOrientation(pick_id)
538
  position = np.float32(pose[0])
539
  return position
540
 
541
  def get_bounding_box(self, obj_name):
542
  obj_id = self.get_obj_id(obj_name)
543
+ return self._p.getAABB(obj_id)
544
 
545
 
546
  class LMP_wrapper():