text: import time
  import os
  import gym
  import gym_panda
  import reflexxes
  import pybullet as p
  import math
  import numpy as np
  import cv2
  import pandas as pd
  import threading
  
  class MovementData:
      def __init__(self, id):
          self.mov_id = id
          self.currentPosition = [0.0, 0.0, 0.0] * 2
          self.currentVelocity = [0.0, 0.0, 0.0] * 2
          self.currentAcceleration = [0.0, 0.0, 0.0] * 2
          self.targetPosition = [0.0, 0.0, 0.0] * 2
          self.targetVelocity = [0.0, 0.0, 0.0] * 2
          self.min_sync_time = 0.0  # in rml, min_sync_time <= time_reach_target_pos_and_vel; if min_sync_time is larger than required, then motion will be less greedy and take time of min_sync; Otherwise, if required time is larger than given min_sync, required time is taken to ensure reach target
  
  
  class Agent:
      def __init__(self, hz=240):
          self.gen = reflexxes.extra.PositionTrajectoryGenerator(
              number_of_dofs=3,
              cycle_time=1/float(hz),
              max_velocity=[5.0, 5.0, 5.0],
              max_acceleration=[10.0, 10.0, 10.0],
              max_jerk=[20.0, 20.0, 20.0])
          self.iter = 0
      def get_obj_position(self, obj_id):
          return p.getBasePositionAndOrientation(obj_id)[0]
  
      def get_obj_orientation(self, obj_id):
          return p.getBasePositionAndOrientation(obj_id)[1]
  
      def get_tip_position(self, env):
          return p.getLinkState(env.pandaUid, 11)[0]
  
      def get_tip_orientation(self, env):
          return p.getLinkState(env.pandaUid, 11)[1]
  
      def gen_motion_list(self, motion_data):
          self.gen.current_position = motion_data.currentPosition[:]
          self.gen.current_velocity = motion_data.currentVelocity[:]
          self.gen.current_acceleration = motion_data.currentAcceleration[:]
  
          pos_list = [motion_data.currentPosition[:]]
          vel_list = [motion_data.currentVelocity[:]]
          acc_list = [motion_data.currentAcceleration[:]]
  
          # generate trajectory
          # gen.trajectory(target_pos, target_vel, min_sync_time)
          for pos, vel, acc in self.gen.trajectory(motion_data.targetPosition, motion_data.targetVelocity, motion_data.min_sync_time):
              pos_list.append(pos)
              vel_list.append(vel)
              acc_list.append(acc)
  
          return pos_list, vel_list, acc_list
  
      def recording(self, env):
          if self.iter == 0:
              if not os.path.exists(env.storage_folder+"/"+env.object+"/"):
                  os.makedirs(env.storage_folder + "/" + env.object + "/")
                  os.makedirs(env.storage_folder + "/" + env.object + "/" + "color/")
                  os.makedirs(env.storage_folder + "/" + env.object + "/" + "cad/")
                  os.makedirs(env.storage_folder + "/" + env.object + "/" + "depth/")
                  os.makedirs(env.storage_folder + "/" + env.object + "/" + "annotations/")
                  os.makedirs(env.storage_folder + "/" + env.object + "/" + "mask/")
              self.folder = env.storage_folder + "/" + env.object + "/"
              self.filenames = []
              self.LinkPositions = []
              self.LinkOrientations = []
              self.LinkR = []
              self.ObjectPositions = []
              self.ObjectOrientations = []
              self.ObjectR = []
  
          self.filenames.append(self.iter)
          self.LinkPositions.append(self.get_tip_position(env))  # recording end effector postion
          self.LinkOrientations.append(self.get_tip_orientation(env))  # recording end effector orientation
          self.LinkR.append(p.getMatrixFromQuaternion(self.LinkOrientations[-1]))
          self.ObjectPositions.append(self.get_obj_position(env.objectUid))
          self.ObjectOrientations.append(self.get_obj_orientation(env.objectUid))
          self.ObjectR.append(p.getMatrixFromQuaternion(self.ObjectOrientations[-1]))
          rgb_filename = self.folder + 'color/%s.jpg' % str(self.iter)
          cad_filename = self.folder + 'cad/%s.jpg' % str(self.iter)
          depth_filename = self.folder + 'depth/%s.png' % str(self.iter)
          annotation_filename = self.folder + 'annotations/%s.png' % str(self.iter)
          mask_filename = self.folder + 'mask/%s.png' % str(self.iter)
          rgb, cad, depth, annotation, mask = env.storage()
          rgb = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
          cad = cv2.cvtColor(cad, cv2.COLOR_RGB2BGR)
          cv2.imwrite(rgb_filename, rgb)
          cv2.imwrite(cad_filename, cad)
          cv2.imwrite(depth_filename, depth)
          cv2.imwrite(annotation_filename, annotation)
          cv2.imwrite(mask_filename, mask)
          self.iter += 1
  
          if env.record_end:
              robot_joints = pd.DataFrame({'filenames': self.filenames,
                                           'LinkPositions': self.LinkPositions,
                                           'LinkOrientations': self.LinkOrientations,
                                           'LinkRotationMatrices': self.LinkR,
                                           'ObjectPositions': self.ObjectPositions,
                                           'ObjectOrientations': self.ObjectOrientations,
                                           'ObjectRotationMatrices': self.ObjectR
                                           })
              robot_joints.to_csv(self.folder + '/robot_joints.csv', index=False)
  
  if __name__ == "__main__":
  
      RECORD = False
      env = gym.make('panda-v0').env
  
      # storage path
      env.storage_folder = os.path.join(os.path.abspath(os.path.dirname(os.getcwd())), "3d_object_reconstruction",
                                         "Data_stuck")
      # prior: grasping position offset w.r.t center of mass of object
      grasp_offset_dict = {
          "YcbPottedMeatCan": [0, 0.005, 0.015],
          "YcbGelatinBox": [0, +0.003, 0.022],
          "YcbMustardBottle": [0, 0, 0.08],
          "YcbTomatoSoupCan": [0, 0.007, 0.025],
          "YcbCrackerBox": [0, -0.01, 0.035],
          "YcbSugarBox": [0, 0, 0.0],
          "YcbBanana": [0, 0, 0],
          "YcbTennisBall": [0, 0, 0.],
      }
      # object to be grasped
      for key in ['YcbTomatoSoupCan', 'YcbCrackerBox']:
      # for key in grasp_offset_dict.keys():
          env.object = key
  
          agent = Agent()
          grasp_offset = grasp_offset_dict[env.object]
  
          env.reset()
          while not env.is_static():
              p.stepSimulation()
  
          debug = True
          if debug:
              p.addUserDebugLine([0, 0, 0], [0.5, 0, 0], [1, 0, 0], lineWidth=5,
                                 parentObjectUniqueId=env.pandaUid, parentLinkIndex=10)
              p.addUserDebugLine([0, 0, 0], [0, 0.5, 0], [0, 1, 0], lineWidth=5,
                                 parentObjectUniqueId=env.pandaUid, parentLinkIndex=10)
              p.addUserDebugLine([0, 0, 0], [0, 0, 0.5], [0, 0, 1], lineWidth=5,
                                 parentObjectUniqueId=env.pandaUid, parentLinkIndex=10)
              p.addUserDebugLine([0, 0, 0], [0.5, 0, 0], [1, 0, 0], lineWidth=5,
                                 parentObjectUniqueId=env.objectUid, parentLinkIndex=-1)
              p.addUserDebugLine([0, 0, 0], [0, 0.5, 0], [0, 1, 0], lineWidth=5,
                                 parentObjectUniqueId=env.objectUid, parentLinkIndex=-1)
              p.addUserDebugLine([0, 0, 0], [0, 0, 0.5], [0, 0, 1], lineWidth=5,
                                 parentObjectUniqueId=env.objectUid, parentLinkIndex=-1)
          fingers = 1
          obj_position = agent.get_obj_position(env.objectUid)
          prepick_position = [x + y for x, y in zip(obj_position, [0, 0, 0.15])]
          grasp_position = [x + y for x, y in zip(obj_position, grasp_offset)]
          init_tip_pose = agent.get_tip_position(env)
          init_tip_ori = agent.get_tip_orientation(env)
          prepick_data = MovementData('prepick')
          prepick_data.currentPosition = init_tip_pose
          prepick_data.targetPosition = prepick_position
          prepick_data.min_sync_time = 2.
  
          pick_data = MovementData('pick')
          pick_data.currentPosition = prepick_position
          pick_data.targetPosition = grasp_position
          pick_data.min_sync_time = 1
  
          prepick_pos, prepick_vel, prepick_acc = agent.gen_motion_list(prepick_data)
          pick_pos, pick_vel, pick_acc = agent.gen_motion_list(pick_data)
  
          pick_group_pos = prepick_pos[:-1] + pick_pos  # remove duplicated pos
          pick_group_vel = prepick_vel[:-1] + pick_vel  # remove duplicated vel
          pick_group_acc = prepick_acc[:-1] + pick_acc  # remove possible duplicated acc to match size
          pick_group_traj_time = np.linspace(0, agent.gen.cycle_time * len(pick_group_pos),
                                             len(pick_group_pos)).tolist()
          pick_time = len(prepick_pos) * agent.gen.cycle_time
  
          # Compute gripper orientation and rotation increments
  
          fingers = [0.1, 0.1]
          for i in range(len(pick_group_pos)):
              action = ['pick', pick_group_pos[i], init_tip_ori, fingers]
              observation, reward, done = env.step(action)
  
          fingers = env.activate()
  
          lift_position = [x + y for x, y in zip(grasp_position, [0, 0, 0.4])]
          lift_data = MovementData('lift')
          lift_data.currentPosition = grasp_position
          lift_data.targetPosition = lift_position
          lift_data.min_sync_time = 3
          lift_pos, lift_vel, lift_acc = agent.gen_motion_list(lift_data)
          ee_position = agent.get_tip_position(env)
          obj_position = agent.get_obj_position(env.objectUid)
          grasp_offset = [x - y for x, y in zip(ee_position, obj_position)]
          for i in range(len(lift_pos)):
              action = ['lift', lift_pos[i], init_tip_ori, fingers]
              observation, reward, done = env.step(action)
          fingers = observation[-2:]
  
  
          # start rotating and recording
          env.init_tip_pose = p.getLinkState(env.pandaUid, 11)[0]
          env.init_tip_ori = p.getLinkState(env.pandaUid, 11)[1]
          env.init_obj_pose = p.getBasePositionAndOrientation(env.objectUid)[0]
          env.init_obj_ori = p.getBasePositionAndOrientation(env.objectUid)[1]
          ee_position = agent.get_tip_position(env)
          # enableCollision = 0
          # p.setCollisionFilterPair(bodyUniqueIdA=env.pandaUid, bodyUniqueIdB=env.objectUid,
          #                          linkIndexA=9, linkIndexB=-1, enableCollision=enableCollision)
          # p.setCollisionFilterPair(bodyUniqueIdA=env.pandaUid, bodyUniqueIdB=env.objectUid,
          #                          linkIndexA=10, linkIndexB=-1, enableCollision=enableCollision)
          for j in range(6):
              rotate_pos = [ee_position] * 120  # hz=240, lasting 0.5 second.
              rotate_group_ori = []
              pre_ori = agent.get_tip_orientation(env)
              for i in range(len(rotate_pos)):
                  rotate_group_ori.append(p.getQuaternionFromEuler(
                      [0., -np.pi, np.pi / 2. + np.pi / 4 * i / len(rotate_pos) + np.pi / 4 * j]
                  ))  # quaternion
              for i in range(len(rotate_pos)):
                  action = ['rotate', rotate_pos[i], rotate_group_ori[i], fingers]
                  observation, reward, done = env.step(action)
                  env.stick_simulation()
                  if RECORD: agent.recording(env)
  
          for j in range(2):
              rotate_pos = [ee_position] * 120  # hz=240, lasting 0.5 second.
              rotate_group_ori = []
              pre_ori = agent.get_tip_orientation(env)
              for i in range(len(rotate_pos)):
                  rotate_group_ori.append(p.getQuaternionFromEuler(
                      [0., -np.pi, 2 * np.pi - np.pi / 4 * i / len(rotate_pos) - np.pi / 4 * j]
                  ))  # quaternion
              for i in range(len(rotate_pos)):
                  action = ['rotate', rotate_pos[i], rotate_group_ori[i], fingers]
                  observation, reward, done = env.step(action)
                  env.stick_simulation()
                  if RECORD: agent.recording(env)
  
          offset = [
              [0.1, 0, 0.1 * (np.sqrt(2) - 1)],
              [0.1 * np.sqrt(2), 0, 0.15 * np.sqrt(2)]
          ]
          rotate_position = ee_position
          for j in range(2):
              pre_position = rotate_position
              rotate_position = [x + y for x, y in zip(ee_position, offset[j])]
              rotate_data = MovementData('rotate')
              rotate_data.currentPosition = pre_position
              rotate_data.targetPosition = rotate_position
              rotate_data.min_sync_time = 0.5
              rotate_pos, rotate_vel, rotate_acc = agent.gen_motion_list(rotate_data)
  
              rotate_group_ori = []
              pre_ori = agent.get_tip_orientation(env)
              for i in range(len(rotate_pos)):
                  rotate_group_ori.append(p.multiplyTransforms(positionA=[0, 0, 0], orientationA=p.getQuaternionFromEuler(
                      [0, -np.pi / 4 * i / len(rotate_pos), 0.]
                  ), positionB=[0, 0, 0], orientationB=pre_ori)[1]
                                          )  # quaternion
              for i in range(len(rotate_pos)):
                  action = ['rotate', rotate_pos[i], rotate_group_ori[i], fingers]
                  observation, reward, done = env.step(action)
                  env.stick_simulation()
                  if RECORD: agent.recording(env)
          action = []
          env.record_end = True
          if RECORD: agent.recording(env)
  
      env.close()
  