correctly handling state

app.py

@@ -16,7 +16,7 @@ from md_logger import MarkdownLogger
 
 
 class DemoRunner:
-
+        
     def __init__(self):
         self._cfg = OmegaConf.to_container(OmegaConf.load('cfg.yaml'), resolve=True)
         self._env = None
@@ -69,12 +69,6 @@ class DemoRunner:
         return lmp_tabletop_ui
 
     def setup(self, api_key, model_name, n_blocks, n_bowls):
-        if not api_key:
-            return 'Please enter your OpenAI API key!', None
-        
-        if n_blocks + n_bowls == 0:
-            return 'Please select at least one object!', None
-
         openai.api_key = api_key
         self._model_name = model_name
 
@@ -112,16 +106,32 @@ class DemoRunner:
         return self._md_logger.get_log(), self._env.get_camera_image(), video_file_name
 
 
-if __name__ == '__main__':
+def setup(api_key, model_name, n_blocks, n_bowls):
+    if not api_key:
+        return 'Please enter your OpenAI API key!', None
+    
+    if n_blocks + n_bowls == 0:
+        return 'Please select at least one object!', None
+
     demo_runner = DemoRunner()
-    demo = gr.Blocks()
+    
+    info, img = demo_runner.setup(api_key, model_name, n_blocks, n_bowls)
+    return info, img, demo_runner
 
+
+def run(instruction, demo_runner):
+    return demo_runner.run(instruction)
+
+
+if __name__ == '__main__':
     with open('README.md', 'r') as f:
         for _ in range(12):
             next(f)
         readme_text = f.read()
 
-    with demo:
+    with gr.Blocks() as demo:
+        state = gr.State(None)
+
         gr.Markdown(readme_text)
         gr.Markdown('# Interactive Demo')
         with gr.Row():
@@ -148,13 +158,13 @@ if __name__ == '__main__':
                 video_run = gr.Video(label='Video of Last Instruction')
         
         btn_setup.click(
-            demo_runner.setup, 
+            setup, 
             inputs=[inp_api_key, inp_model_name, inp_n_blocks, inp_n_bowls], 
-            outputs=[info_setup, img_setup]
+            outputs=[info_setup, img_setup, state]
         )
         btn_run.click(
-            demo_runner.run, 
-            inputs=[inp_instruction], 
+            run, 
+            inputs=[inp_instruction, state], 
             outputs=[info_run, img_setup, video_run]
         )
         

sim.py

@@ -1,4 +1,5 @@
 import pybullet
+from pybullet_utils.bullet_client import BulletClient
 import pybullet_data
 import threading
 from time import sleep
@@ -7,26 +8,28 @@ import os
 from consts import BOUNDS, COLORS, PIXEL_SIZE, CORNER_POS
 from shapely.geometry import box
 
+
 # Gripper (Robotiq 2F85) code
 class Robotiq2F85:
   """Gripper handling for Robotiq 2F85."""
 
-  def __init__(self, robot, tool):
+  def __init__(self, robot, tool, p):
     self.robot = robot
     self.tool = tool
+    self._p = p
     pos = [0.1339999999999999, -0.49199999999872496, 0.5]
-    rot = pybullet.getQuaternionFromEuler([np.pi, 0, np.pi])
+    rot = self._p.getQuaternionFromEuler([np.pi, 0, np.pi])
     urdf = 'robotiq_2f_85/robotiq_2f_85.urdf'
-    self.body = pybullet.loadURDF(urdf, pos, rot)
-    self.n_joints = pybullet.getNumJoints(self.body)
+    self.body = self._p.loadURDF(urdf, pos, rot)
+    self.n_joints = self._p.getNumJoints(self.body)
     self.activated = False
 
     # Connect gripper base to robot tool.
-    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]))
+    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]))
 
     # Set friction coefficients for gripper fingers.
-    for i in range(pybullet.getNumJoints(self.body)):
-      pybullet.changeDynamics(self.body, i, lateralFriction=10.0, spinningFriction=1.0, rollingFriction=1.0, frictionAnchor=True)
+    for i in range(self._p.getNumJoints(self.body)):
+      self._p.changeDynamics(self.body, i, lateralFriction=10.0, spinningFriction=1.0, rollingFriction=1.0, frictionAnchor=True)
 
     # Start thread to handle additional gripper constraints.
     self.motor_joint = 1
@@ -39,22 +42,22 @@ class Robotiq2F85:
   def step(self):
     while True:
       try:
-        currj = [pybullet.getJointState(self.body, i)[0] for i in range(self.n_joints)]
+        currj = [self._p.getJointState(self.body, i)[0] for i in range(self.n_joints)]
         indj = [6, 3, 8, 5, 10]
         targj = [currj[1], -currj[1], -currj[1], currj[1], currj[1]]
-        pybullet.setJointMotorControlArray(self.body, indj, pybullet.POSITION_CONTROL, targj, positionGains=np.ones(5))
+        self._p.setJointMotorControlArray(self.body, indj, self._p.POSITION_CONTROL, targj, positionGains=np.ones(5))
       except:
         return
       sleep(0.001)
 
   # Close gripper fingers.
   def activate(self):
-    pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=1, force=10)
+    self._p.setJointMotorControl2(self.body, self.motor_joint, self._p.VELOCITY_CONTROL, targetVelocity=1, force=10)
     self.activated = True
 
   # Open gripper fingers.
   def release(self):
-    pybullet.setJointMotorControl2(self.body, self.motor_joint, pybullet.VELOCITY_CONTROL, targetVelocity=-1, force=10)
+    self._p.setJointMotorControl2(self.body, self.motor_joint, self._p.VELOCITY_CONTROL, targetVelocity=-1, force=10)
     self.activated = False
 
   # If activated and object in gripper: check object contact.
@@ -75,7 +78,7 @@ class Robotiq2F85:
   def external_contact(self, body=None):
     if body is None:
       body = self.body
-    pts = pybullet.getContactPoints(bodyA=body)
+    pts = self._p.getContactPoints(bodyA=body)
     pts = [pt for pt in pts if pt[2] != self.body]
     return len(pts) > 0  # pylint: disable=g-explicit-length-test
 
@@ -86,17 +89,17 @@ class Robotiq2F85:
     return success
 
   def grasp_width(self):
-    lpad = np.array(pybullet.getLinkState(self.body, 4)[0])
-    rpad = np.array(pybullet.getLinkState(self.body, 9)[0])
+    lpad = np.array(self._p.getLinkState(self.body, 4)[0])
+    rpad = np.array(self._p.getLinkState(self.body, 9)[0])
     dist = np.linalg.norm(lpad - rpad) - 0.047813
     return dist
 
   def check_proximity(self):
-    ee_pos = np.array(pybullet.getLinkState(self.robot, self.tool)[0])
-    tool_pos = np.array(pybullet.getLinkState(self.body, 0)[0])
+    ee_pos = np.array(self._p.getLinkState(self.robot, self.tool)[0])
+    tool_pos = np.array(self._p.getLinkState(self.body, 0)[0])
     vec = (tool_pos - ee_pos) / np.linalg.norm((tool_pos - ee_pos))
     ee_targ = ee_pos + vec
-    ray_data = pybullet.rayTest(ee_pos, ee_targ)[0]
+    ray_data = self._p.rayTest(ee_pos, ee_targ)[0]
     obj, link, ray_frac = ray_data[0], ray_data[1], ray_data[2]
     return obj, link, ray_frac
 
@@ -108,16 +111,15 @@ class PickPlaceEnv():
     self.dt = 1/480
     self.sim_step = 0
 
-    # Configure and start PyBullet.
-    # python3 -m pybullet_utils.runServer
-    # pybullet.connect(pybullet.SHARED_MEMORY)  # pybullet.GUI for local GUI.
-    pybullet.connect(pybullet.DIRECT)  # pybullet.GUI for local GUI.
-    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_GUI, 0)
-    pybullet.setPhysicsEngineParameter(enableFileCaching=0)
+    # Configure and start PyBullet
+    # self._p = pybullet.connect(pybullet.DIRECT)
+    self._p = BulletClient(connection_mode=pybullet.DIRECT)
+    self._p.configureDebugVisualizer(self._p.COV_ENABLE_GUI, 0)
+    self._p.setPhysicsEngineParameter(enableFileCaching=0)
     assets_path = os.path.dirname(os.path.abspath(""))
-    pybullet.setAdditionalSearchPath(assets_path)
-    pybullet.setAdditionalSearchPath(pybullet_data.getDataPath())
-    pybullet.setTimeStep(self.dt)
+    self._p.setAdditionalSearchPath(assets_path)
+    self._p.setAdditionalSearchPath(pybullet_data.getDataPath())
+    self._p.setTimeStep(self.dt)
 
     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).
     self.home_ee_euler = (np.pi, 0, np.pi)  # (RX, RY, RZ) rotation in Euler angles.
@@ -130,36 +132,36 @@ class PickPlaceEnv():
     self.high_frame_rate = high_frame_rate
 
   def reset(self, object_list):
-    pybullet.resetSimulation(pybullet.RESET_USE_DEFORMABLE_WORLD)
-    pybullet.setGravity(0, 0, -9.8)
+    self._p.resetSimulation(self._p.RESET_USE_DEFORMABLE_WORLD)
+    self._p.setGravity(0, 0, -9.8)
     self.cache_video = []
 
     # Temporarily disable rendering to load URDFs faster.
-    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 0)
+    self._p.configureDebugVisualizer(self._p.COV_ENABLE_RENDERING, 0)
 
     # Add robot.
-    pybullet.loadURDF("plane.urdf", [0, 0, -0.001])
-    self.robot_id = pybullet.loadURDF("ur5e/ur5e.urdf", [0, 0, 0], flags=pybullet.URDF_USE_MATERIAL_COLORS_FROM_MTL)
-    self.ghost_id = pybullet.loadURDF("ur5e/ur5e.urdf", [0, 0, -10])  # For forward kinematics.
-    self.joint_ids = [pybullet.getJointInfo(self.robot_id, i) for i in range(pybullet.getNumJoints(self.robot_id))]
-    self.joint_ids = [j[0] for j in self.joint_ids if j[2] == pybullet.JOINT_REVOLUTE]
+    self._p.loadURDF("plane.urdf", [0, 0, -0.001])
+    self.robot_id = self._p.loadURDF("ur5e/ur5e.urdf", [0, 0, 0], flags=self._p.URDF_USE_MATERIAL_COLORS_FROM_MTL)
+    self.ghost_id = self._p.loadURDF("ur5e/ur5e.urdf", [0, 0, -10])  # For forward kinematics.
+    self.joint_ids = [self._p.getJointInfo(self.robot_id, i) for i in range(self._p.getNumJoints(self.robot_id))]
+    self.joint_ids = [j[0] for j in self.joint_ids if j[2] == self._p.JOINT_REVOLUTE]
 
     # Move robot to home configuration.
     for i in range(len(self.joint_ids)):
-      pybullet.resetJointState(self.robot_id, self.joint_ids[i], self.home_joints[i])
+      self._p.resetJointState(self.robot_id, self.joint_ids[i], self.home_joints[i])
 
     # Add gripper.
     if self.gripper is not None:
       while self.gripper.constraints_thread.is_alive():
         self.constraints_thread_active = False
-    self.gripper = Robotiq2F85(self.robot_id, self.ee_link_id)
+    self.gripper = Robotiq2F85(self.robot_id, self.ee_link_id, self._p)
     self.gripper.release()
 
     # Add workspace.
-    plane_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])
-    plane_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])
-    plane_id = pybullet.createMultiBody(0, plane_shape, plane_visual, basePosition=[0, -0.5, 0])
-    pybullet.changeVisualShape(plane_id, -1, rgbaColor=[0.2, 0.2, 0.2, 1.0])
+    plane_shape = self._p.createCollisionShape(self._p.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])
+    plane_visual = self._p.createVisualShape(self._p.GEOM_BOX, halfExtents=[0.3, 0.3, 0.001])
+    plane_id = self._p.createMultiBody(0, plane_shape, plane_visual, basePosition=[0, -0.5, 0])
+    self._p.changeVisualShape(plane_id, -1, rgbaColor=[0.2, 0.2, 0.2, 1.0])
 
     # Load objects according to config.
     self.object_list = object_list
@@ -186,21 +188,20 @@ class PickPlaceEnv():
         object_type = obj_name.split(' ')[1]
         object_position = rand_xyz.squeeze()
         if object_type == 'block':
-          object_shape = pybullet.createCollisionShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])
-          object_visual = pybullet.createVisualShape(pybullet.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])
-          object_id = pybullet.createMultiBody(0.01, object_shape, object_visual, basePosition=object_position)
+          object_shape = self._p.createCollisionShape(self._p.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])
+          object_visual = self._p.createVisualShape(self._p.GEOM_BOX, halfExtents=[0.02, 0.02, 0.02])
+          object_id = self._p.createMultiBody(0.01, object_shape, object_visual, basePosition=object_position)
         elif object_type == 'bowl':
           object_position[2] = 0
-          object_id = pybullet.loadURDF("bowl/bowl.urdf", object_position, useFixedBase=1)
-        pybullet.changeVisualShape(object_id, -1, rgbaColor=object_color)
+          object_id = self._p.loadURDF("bowl/bowl.urdf", object_position, useFixedBase=1)
+        self._p.changeVisualShape(object_id, -1, rgbaColor=object_color)
         self.obj_name_to_id[obj_name] = object_id
-    
 
     # Re-enable rendering.
-    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_RENDERING, 1)
+    self._p.configureDebugVisualizer(self._p.COV_ENABLE_RENDERING, 1)
 
     for _ in range(200):
-      pybullet.stepSimulation()
+      self._p.stepSimulation()
 
     # record object positions at reset
     self.init_pos = {name: self.get_obj_pos(name) for name in object_list}
@@ -209,25 +210,25 @@ class PickPlaceEnv():
 
   def servoj(self, joints):
     """Move to target joint positions with position control."""
-    pybullet.setJointMotorControlArray(
+    self._p.setJointMotorControlArray(
       bodyIndex=self.robot_id,
       jointIndices=self.joint_ids,
-      controlMode=pybullet.POSITION_CONTROL,
+      controlMode=self._p.POSITION_CONTROL,
       targetPositions=joints,
       positionGains=[0.01]*6)
   
   def movep(self, position):
     """Move to target end effector position."""
-    joints = pybullet.calculateInverseKinematics(
+    joints = self._p.calculateInverseKinematics(
         bodyUniqueId=self.robot_id,
         endEffectorLinkIndex=self.tip_link_id,
         targetPosition=position,
-        targetOrientation=pybullet.getQuaternionFromEuler(self.home_ee_euler),
+        targetOrientation=self._p.getQuaternionFromEuler(self.home_ee_euler),
         maxNumIterations=100)
     self.servoj(joints)
 
   def get_ee_pos(self):
-    ee_xyz = np.float32(pybullet.getLinkState(self.robot_id, self.tip_link_id)[0])
+    ee_xyz = np.float32(self._p.getLinkState(self.robot_id, self.tip_link_id)[0])
     return ee_xyz
 
   def step(self, action=None):
@@ -306,17 +307,17 @@ class PickPlaceEnv():
 
   def set_alpha_transparency(self, alpha: float) -> None:
     for id in range(20):
-      visual_shape_data = pybullet.getVisualShapeData(id)
+      visual_shape_data = self._p.getVisualShapeData(id)
       for i in range(len(visual_shape_data)):
         object_id, link_index, _, _, _, _, _, rgba_color = visual_shape_data[i]
         rgba_color = list(rgba_color[0:3]) +  [alpha]
-        pybullet.changeVisualShape(
+        self._p.changeVisualShape(
             self.robot_id, linkIndex=i, rgbaColor=rgba_color)      
-        pybullet.changeVisualShape(
+        self._p.changeVisualShape(
             self.gripper.body, linkIndex=i, rgbaColor=rgba_color)
 
   def step_sim_and_render(self):
-    pybullet.stepSimulation()
+    self._p.stepSimulation()
     self.sim_step += 1
 
     interval = 40 if self.high_frame_rate else 60
@@ -346,7 +347,7 @@ class PickPlaceEnv():
     # Get heightmaps and colormaps.
     points = self.get_pointcloud(depth, intrinsics)
     position = np.float32(position).reshape(3, 1)
-    rotation = pybullet.getMatrixFromQuaternion(orientation)
+    rotation = self._p.getMatrixFromQuaternion(orientation)
     rotation = np.float32(rotation).reshape(3, 3)
     transform = np.eye(4)
     transform[:3, :] = np.hstack((rotation, position))
@@ -363,37 +364,37 @@ class PickPlaceEnv():
     # Camera parameters.
     position = (0, -0.85, 0.4)
     orientation = (np.pi / 4 + np.pi / 48, np.pi, np.pi)
-    orientation = pybullet.getQuaternionFromEuler(orientation)
+    orientation = self._p.getQuaternionFromEuler(orientation)
     zrange = (0.01, 10.)
     noise=True
 
     # OpenGL camera settings.
     lookdir = np.float32([0, 0, 1]).reshape(3, 1)
     updir = np.float32([0, -1, 0]).reshape(3, 1)
-    rotation = pybullet.getMatrixFromQuaternion(orientation)
+    rotation = self._p.getMatrixFromQuaternion(orientation)
     rotm = np.float32(rotation).reshape(3, 3)
     lookdir = (rotm @ lookdir).reshape(-1)
     updir = (rotm @ updir).reshape(-1)
     lookat = position + lookdir
     focal_len = intrinsics[0]
     znear, zfar = (0.01, 10.)
-    viewm = pybullet.computeViewMatrix(position, lookat, updir)
+    viewm = self._p.computeViewMatrix(position, lookat, updir)
     fovh = (image_size[0] / 2) / focal_len
     fovh = 180 * np.arctan(fovh) * 2 / np.pi
 
     # Notes: 1) FOV is vertical FOV 2) aspect must be float
     aspect_ratio = image_size[1] / image_size[0]
-    projm = pybullet.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)
+    projm = self._p.computeProjectionMatrixFOV(fovh, aspect_ratio, znear, zfar)
 
     # Render with OpenGL camera settings.
-    _, _, color, depth, segm = pybullet.getCameraImage(
+    _, _, color, depth, segm = self._p.getCameraImage(
         width=image_size[1],
         height=image_size[0],
         viewMatrix=viewm,
         projectionMatrix=projm,
         shadow=1,
-        flags=pybullet.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,
-        renderer=pybullet.ER_BULLET_HARDWARE_OPENGL)
+        flags=self._p.ER_SEGMENTATION_MASK_OBJECT_AND_LINKINDEX,
+        renderer=self._p.ER_BULLET_HARDWARE_OPENGL)
 
     # Get color image.
     color_image_size = (image_size[0], image_size[1], 4)
@@ -533,13 +534,13 @@ class PickPlaceEnv():
       position = np.float32(np.array(CORNER_POS[obj_name]))
     else:
       pick_id = self.get_obj_id(obj_name)
-      pose = pybullet.getBasePositionAndOrientation(pick_id)
+      pose = self._p.getBasePositionAndOrientation(pick_id)
       position = np.float32(pose[0])
     return position
   
   def get_bounding_box(self, obj_name):
     obj_id = self.get_obj_id(obj_name)
-    return pybullet.getAABB(obj_id)
+    return self._p.getAABB(obj_id)
 
 
 class LMP_wrapper():