class TeleopEnv(): def __init__(self): # create simulation (GUI) self.urdfRootPath = pybullet_data.getDataPath() p.connect(p.GUI) p.setGravity(0, 0, -9.81) # set up camera self._set_camera() # load some scene objects p.loadURDF(os.path.join(self.urdfRootPath, "plane.urdf"), basePosition=[0, 0, -0.65]) p.loadURDF(os.path.join(self.urdfRootPath, "table/table.urdf"), basePosition=[0.5, 0, -0.65]) # example YCB object obj1 = YCBObject('003_cracker_box') obj1.load() p.resetBasePositionAndOrientation(obj1.body_id, [0.7, -0.2, 0.1], [0, 0, 0, 1]) # load a panda robot self.panda = Panda() # connect to joystick for control self.joystick = Joystick(scale=0.01) def reset(self): self.panda.reset() return self.panda.state def close(self): p.disconnect() def step(self): # get current state state = self.panda.state # read in from joystick input = self.joystick.get_controller_state() dpos, dquat, grasp, reset = ( input["dpos"], input["dquat"], input["grasp"], input["reset"], ) # action in this example is the end-effector velocity self.panda.step(dposition=dpos, dquaternion=dquat, grasp_open=grasp) # take simulation step p.stepSimulation() # return next_state, reward, done, info next_state = self.panda.state reward = 0.0 done = False if reset: done = True info = {} return next_state, reward, done, info def render(self): (width, height, pxl, depth, segmentation) = p.getCameraImage(width=self.camera_width, height=self.camera_height, viewMatrix=self.view_matrix, projectionMatrix=self.proj_matrix) rgb_array = np.array(pxl, dtype=np.uint8) rgb_array = np.reshape(rgb_array, (self.camera_height, self.camera_width, 4)) rgb_array = rgb_array[:, :, :3] return rgb_array def _set_camera(self): self.camera_width = 256 self.camera_height = 256 p.resetDebugVisualizerCamera(cameraDistance=1.2, cameraYaw=30, cameraPitch=-60, cameraTargetPosition=[0.5, -0.2, 0.0]) self.view_matrix = p.computeViewMatrixFromYawPitchRoll( cameraTargetPosition=[0.5, 0, 0], distance=1.0, yaw=90, pitch=-50, roll=0, upAxisIndex=2) self.proj_matrix = p.computeProjectionMatrixFOV( fov=60, aspect=float(self.camera_width) / self.camera_height, nearVal=0.1, farVal=100.0)