class RattlerCamera(object): def __init__(self): self._head_camera = VisionSensor('rattler_eye') self._head_camera_mask = None self._head_camera_state = 1 def get_camera_state(self) -> int: return self._head_camera_state def set_camera_state(self, state: int) -> None: self._head_camera_state = state def set_camera_properties(self, conf: CameraConfig) -> None: if not (conf.rgb or conf.depth): self._head_camera.remove() else: self._head_camera.set_resolution(conf.image_size) self._head_camera.set_render_mode(conf.render_mode) if conf.mask: self._head_camera_mask = self._head_camera.copy() self._head_camera_mask.set_render_mode( RenderMode.OPENGL_COLOR_CODED) def capture_rgb(self) -> np.ndarray: return self._head_camera.capture_rgb() def capture_depth(self) -> np.ndarray: return self._head_camera.capture_depth() def mask(self) -> VisionSensor: return self._head_camera_mask
class ArmECM(PyRep): def __init__(self, pr): #super(ArmECM, self).__init__(scenePath) """self.pr = PyRep() self.pr.launch(scenePath) self.pr.start() self.pr.step()""" self.pr = pr self.base_handle = Shape('L0_respondable_ECM') self.j1_handle = Joint('J1_ECM') self.j2_handle = Joint('J2_ECM') self.j3_handle = Joint('J3_ECM') self.j4_handle = Joint('J4_ECM') self.left_cam = VisionSensor('Vision_sensor_left') self.right_cam = VisionSensor('Vision_sensor_right') def getJointAngles(self): pos1 = self.j1_handle.get_joint_position() pos2 = self.j2_handle.get_joint_position() pos3 = self.j3_handle.get_joint_position() pos4 = self.j4_handle.get_joint_position() return np.array([pos1, pos2, pos3, pos4]) def setJointAngles(self, jointAngles): self.j1_handle.set_joint_position(jointAngles[0]) self.j2_handle.set_joint_position(jointAngles[1]) self.j3_handle.set_joint_position(jointAngles[2]) self.j4_handle.set_joint_position(jointAngles[3]) def getCurrentPose(self): return self.left_cam.get_pose(relative_to=self.base_handle) def getStereoImagePairs(self): return self.left_cam.capture_rgb(), self.right_cam.capture_rgb() def getDepthImagePairs(self): return self.left_cam.capture_depth(True), self.right_cam.capture_depth( True) """def stopSim(self):
def get_rgb_depth(sensor: VisionSensor, get_rgb: bool, get_depth: bool, rgb_noise: NoiseModel, depth_noise: NoiseModel): rgb = depth = None if sensor is not None and (get_rgb or get_depth): sensor.handle_explicitly() if get_rgb: rgb = sensor.capture_rgb() if rgb_noise is not None: rgb = rgb_noise.apply(rgb) if get_depth: depth = sensor.capture_depth() if depth_noise is not None: depth = depth_noise.apply(depth) return rgb, depth
class TestVisionSensors(TestCore): def setUp(self): super().setUp() [self.pyrep.step() for _ in range(10)] self.cam = VisionSensor('cam0') def test_capture_rgb(self): img = self.cam.capture_rgb() self.assertEqual(img.shape, (16, 16, 3)) # Check that it's not a blank image self.assertFalse(img.min() == img.max() == 0.0) def test_capture_depth(self): img = self.cam.capture_depth() self.assertEqual(img.shape, (16, 16)) # Check that it's not a blank depth map self.assertFalse(img.min() == img.max() == 1.0)
# current_pose = get_2d_pose(robot) target_pose = get_2d_pose(goal) current_pose = robot.get_position(vision_sensor)[:2] + robot.get_orientation()[-1:] current_pose = (-current_pose[0], current_pose[1], current_pose[2]) print(current_pose) print(get_2d_pose(robot)) pr.step() origin_x, _, origin_y, _, _, _ = vision_sensor.get_bounding_box() # Setup occ_grid occ_grid = OccupancyGrid(origin_x, origin_y) depth = vision_sensor.capture_depth() occ_grid.fromDepth(depth) ''' dij_solver = Dijkstra(occ_grid, current_pose, target_pose) moves = dij_solver.solve() for i in moves: pr.step() set_2d_pose(robot, (i[0], i[1], current_pose[2])) time.sleep(0.05) ''' ''' # RRT Solver rrt_solver = RRT(pr, robot, current_pose, target_pose, [-0.5,0.5], [-0.5,0.5]) moves = rrt_solver.solve()
class EnvPollos(Env): def __init__(self, ep_length=100): """ Pollos environment for testing purposes :param dim: (int) the size of the dimensions you want to learn :param ep_length: (int) the length of each episodes in timesteps """ logging.basicConfig(level=logging.DEBUG) # they have to be simmetric. We interpret actions as angular increments #self.action_space = Box(low=np.array([-0.1, -1.7, -2.5, -1.5, 0.0, 0.0]), # high=np.array([0.1, 1.7, 2.5, 1.5, 0.0, 0.0])) inc = 0.1 self.action_space = Box(low=np.array( [-inc, -inc, -inc, -inc, 0.0, 0.0]), high=np.array([inc, inc, inc, inc, 0.0, 0.0])) self.observation_space = Box(low=0, high=255, shape=(480, 640, 1), dtype=np.uint8) # self.observation_space = Box(low=np.array([-0.1, -0.2, 0, -1.5, 0.0, 0.0, -0.5, -0.2, -0.2]), # high=np.array([0.1, 1.7, 2.5, 0, 0.0, 0.0, 0.5, 0.2, 1.0])) # self.pr = PyRep() self.pr.launch(SCENE_FILE, headless=False) self.pr.start() self.agent = UR10() self.agent.max_velocity = 1.2 self.agent.set_control_loop_enabled(True) self.agent.set_motor_locked_at_zero_velocity(True) self.joints = [ Joint('UR10_joint1'), Joint('UR10_joint2'), Joint('UR10_joint3'), Joint('UR10_joint4'), Joint('UR10_joint5'), Joint('UR10_joint6') ] #self.joints_limits = [[j.get_joint_interval()[1][0],j.get_joint_interval()[1][0]+j.get_joint_interval()[1][1]] # for j in self.joints] self.high_joints_limits = [0.1, 1.7, 2.7, 0.0, 0.1, 0.1] self.low_joints_limits = [-0.1, -0.2, 0.0, -1.5, -0.1, -0.1] self.agent_hand = Shape('hand') self.initial_agent_tip_position = self.agent.get_tip().get_position() self.initial_agent_tip_quaternion = self.agent.get_tip( ).get_quaternion() self.target = Dummy('UR10_target') self.initial_joint_positions = self.agent.get_joint_positions() self.pollo_target = Dummy('pollo_target') self.pollo = Shape('pollo') self.table_target = Dummy('table_target') self.initial_pollo_position = self.pollo.get_position() self.initial_pollo_orientation = self.pollo.get_quaternion() self.table_target = Dummy('table_target') # objects self.scene_objects = [ Shape('table0'), Shape('Plane'), Shape('Plane0'), Shape('ConcretBlock') ] #self.agent_tip = self.agent.get_tip() self.initial_distance = np.linalg.norm( np.array(self.initial_pollo_position) - np.array(self.initial_agent_tip_position)) # camera self.camera = VisionSensor('kinect_depth') self.camera_matrix_extrinsics = vrep.simGetObjectMatrix( self.camera.get_handle(), -1) self.np_camera_matrix_extrinsics = np.delete( np.linalg.inv(self.vrepToNP(self.camera_matrix_extrinsics)), 3, 0) width = 640.0 height = 480.0 angle = math.radians(57.0) focalx_px = (width / 2.0) / math.tan(angle / 2.0) focaly_px = (height / 2.0) / math.tan(angle / 2.0) self.np_camera_matrix_intrinsics = np.array([[-focalx_px, 0.0, 320.0], [0.0, -focalx_px, 240.0], [0.0, 0.0, 1.0]]) self.reset() def reset(self): pos = list( np.random.uniform([-0.1, -0.1, 0.0], [0.1, 0.1, 0.1]) + self.initial_pollo_position) self.pollo.set_position(pos) self.pollo.set_quaternion(self.initial_pollo_orientation) self.agent.set_joint_positions(self.initial_joint_positions) self.initial_epoch_time = time.time() while True: # wait for arm to stop self.pr.step() # Step the physics simulation a = self.agent.get_joint_velocities() if not np.any(np.where(np.fabs(a) < 0.1, False, True)): break # a los dos mil reset hay que recargar el ROBOT porque se descoyunta #if self.num_resets > 2000: self.total_reward = 0.0 print('----------------') return self._get_state() def step(self, action): if action is None: self.pr.step() return self._get_state(), 0.0, False, {} # check for nan if np.any(np.isnan(action)): print(action) return self._get_state(), -1.0, False, {} # action[1] = np.interp(action[1], [-1,7,1.7], [-0.2,1.7]) # action[2] = np.interp(action[2], [-2.5,2.5], [0,2.5]) # action[3] = np.interp(action[3], [-1.5,1.5], [-1.5,0]) # add actions as increments to current joints value new_joints = np.array( self.agent.get_joint_positions()) + np.array(action) # check limits if np.any(np.greater(new_joints, self.high_joints_limits)) or np.any( np.less(new_joints, self.low_joints_limits)): logging.debug("New joints value out of limits r=-10") return self._get_state(), -10.0, True, {} # move the robot and wait to stop self.agent.set_joint_target_positions(new_joints) reloj = time.time() while True: # wait for arm to stop self.pr.step() # Step the physics simulation a = self.agent.get_joint_velocities() if not np.any(np.where(np.fabs(a) < 0.1, False, True)) or (time.time() - reloj) > 3: break # compute relevant magnitudes np_pollo_target = np.array(self.pollo_target.get_position()) np_robot_tip_position = np.array(self.agent.get_tip().get_position()) dist = np.linalg.norm(np_pollo_target - np_robot_tip_position) altura = np.clip((np_pollo_target[2] - self.initial_pollo_position[2]), 0, 0.5) for obj in self.scene_objects: # colisión con la mesa if self.agent_hand.check_collision(obj): logging.debug("Collision with objects r=-10") self.total_reward = self.total_reward + -10.0 print(self.total_reward) return self._get_state(), -10.0, True, {} if altura > 0.3: # pollo en mano logging.debug("Height reached !!! r=30") self.total_reward = self.total_reward + 30.0 print(self.total_reward) return self._get_state(), 30.0, True, {} elif dist > self.initial_distance: # mano se aleja logging.debug("Hand moving away from chicken r=-10") self.total_reward = self.total_reward + -10.0 print(self.total_reward) return self._get_state(), -10.0, True, {} used_time = time.time() - self.initial_epoch_time if used_time > 5: # too much time logging.debug("Time out. End of epoch r=-10") if altura > 0.01: self.total_reward = self.total_reward + altura else: self.total_reward = self.total_reward - 10 print(self.total_reward) return self._get_state(), -10.0, True, {} # Reward #pollo_height = np.exp(-altura*20) # para 0.3m pollo_height = 0.002, para 0m pollo_height = 1 reward = altura - dist - used_time self.total_reward = self.total_reward + reward #print(self.total_reward) return self._get_state(), reward, False, {} def close(self): self.pr.stop() self.pr.shutdown() def render(self): print("RENDER") np_pollo_en_camara = self.getPolloEnCamara() # capture depth image depth = self.camera.capture_rgb() circ = plt.Circle( (int(np_pollo_en_camara[0]), int(np_pollo_en_camara[1])), 10) plt.clf() fig = plt.gcf() ax = fig.gca() ax.add_artist(circ) ax.imshow(depth, cmap='hot') plt.pause(0.000001) # Aux # transform env.pollo_target.get_position() to camera coordinates and project pollo_en_camera a image coordinates def getPolloEnCamara(self): np_pollo_target = np.array(self.pollo_target.get_position()) np_pollo_target_cam = self.np_camera_matrix_extrinsics.dot( np.append([np_pollo_target], 1.0)) np_pollo_en_camara = self.np_camera_matrix_intrinsics.dot( np_pollo_target_cam) np_pollo_en_camara = np_pollo_en_camara / np_pollo_en_camara[2] np_pollo_en_camara = np.delete(np_pollo_en_camara, 2) return np_pollo_en_camara def getPolloEnCamaraEx(self): np_pollo_target = np.array(self.pollo_target.get_position()) np_pollo_en_camara = self.np_camera_matrix_extrinsics.dot( np.append([np_pollo_target], 1.0)) return np_pollo_en_camara def _get_state(self): depth = self.camera.capture_depth() scaler = MinMaxScaler(feature_range=(0, 250)) scaler = scaler.fit(depth) depth_scaled = scaler.transform(depth) depth_scaled = np.array(depth_scaled).reshape(480, 640, 1) return depth_scaled.astype('uint8') def vrepToNP(self, c): return np.array([[c[0], c[4], c[8], c[3]], [c[1], c[5], c[9], c[7]], [c[2], c[6], c[10], c[11]], [0, 0, 0, 1]])
class TestVisionSensors(TestCore): def setUp(self): super().setUp() [self.pyrep.step() for _ in range(10)] self.cam = VisionSensor('cam0') def test_handle_explicitly(self): cam = VisionSensor.create((640, 480)) # blank image rgb = cam.capture_rgb() self.assertEqual(rgb.sum(), 0) # non-blank image cam.set_explicit_handling(value=1) cam.handle_explicitly() rgb = cam.capture_rgb() self.assertNotEqual(rgb.sum(), 0) cam.remove() def test_capture_rgb(self): img = self.cam.capture_rgb() self.assertEqual(img.shape, (16, 16, 3)) # Check that it's not a blank image self.assertFalse(img.min() == img.max() == 0.0) def test_capture_depth(self): img = self.cam.capture_depth() self.assertEqual(img.shape, (16, 16)) # Check that it's not a blank depth map self.assertFalse(img.min() == img.max() == 1.0) def test_create(self): cam = VisionSensor.create([640, 480], perspective_mode=True, view_angle=35.0, near_clipping_plane=0.25, far_clipping_plane=5.0, render_mode=RenderMode.OPENGL3) self.assertEqual(cam.capture_rgb().shape, (480, 640, 3)) self.assertEqual(cam.get_perspective_mode(), PerspectiveMode.PERSPECTIVE) self.assertAlmostEqual(cam.get_perspective_angle(), 35.0, 3) self.assertEqual(cam.get_near_clipping_plane(), 0.25) self.assertEqual(cam.get_far_clipping_plane(), 5.0) self.assertEqual(cam.get_render_mode(), RenderMode.OPENGL3) def test_get_set_resolution(self): self.cam.set_resolution([320, 240]) self.assertEqual(self.cam.get_resolution(), [320, 240]) self.assertEqual(self.cam.capture_rgb().shape, (240, 320, 3)) def test_get_set_perspective_mode(self): for perspective_mode in PerspectiveMode: self.cam.set_perspective_mode(perspective_mode) self.assertEqual(self.cam.get_perspective_mode(), perspective_mode) def test_get_set_render_mode(self): for render_mode in [RenderMode.OPENGL, RenderMode.OPENGL3]: self.cam.set_render_mode(render_mode) self.assertEqual(self.cam.get_render_mode(), render_mode) def test_get_set_perspective_angle(self): self.cam.set_perspective_angle(45.0) self.assertAlmostEqual(self.cam.get_perspective_angle(), 45.0, 3) def test_get_set_orthographic_size(self): self.cam.set_orthographic_size(1.0) self.assertEqual(self.cam.get_orthographic_size(), 1.0) def test_get_set_near_clipping_plane(self): self.cam.set_near_clipping_plane(0.5) self.assertEqual(self.cam.get_near_clipping_plane(), 0.5) def test_get_set_far_clipping_plane(self): self.cam.set_far_clipping_plane(0.5) self.assertEqual(self.cam.get_far_clipping_plane(), 0.5) def test_get_set_windowed_size(self): self.cam.set_windowed_size((640, 480)) self.assertEqual(self.cam.get_windowed_size(), (640, 480)) def test_get_set_entity_to_render(self): self.cam.set_entity_to_render(-1) self.assertEqual(self.cam.get_entity_to_render(), -1)
class LoCoBotCamera(Camera): """docstring for SimpleCamera""" def __init__(self, configs, simulator): self.sim = simulator.sim self.rgb_cam = VisionSensor("kinect_rgb") self.depth_cam = VisionSensor("kinect_depth") self.rgb_cam.set_render_mode(RenderMode.OPENGL3) self.depth_cam.set_render_mode(RenderMode.OPENGL3) # Pan and tilt related variables. self.pan_joint = Joint("LoCoBot_head_pan_joint") self.tilt_joint = Joint("LoCoBot_head_tilt_joint") def get_rgb(self): return self.rgb_cam.capture_rgb() def get_depth(self): return self.depth_cam.capture_depth() def get_rgb_depth(self): return self.get_rgb(), self.get_depth() def get_intrinsics(self): # Todo: Remove this after we fix intrinsics raise NotImplementedError """ Returns the instrinsic matrix of the camera :return: the intrinsic matrix (shape: :math:`[3, 3]`) :rtype: np.ndarray """ # fx = self.configs['Camera.fx'] # fy = self.configs['Camera.fy'] # cx = self.configs['Camera.cx'] # cy = self.configs['Camera.cy'] Itc = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) return Itc def pix_to_3dpt(self, rs, cs, in_cam=False): """ Get the 3D points of the pixels in RGB images. :param rs: rows of interest in the RGB image. It can be a list or 1D numpy array which contains the row indices. The default value is None, which means all rows. :param cs: columns of interest in the RGB image. It can be a list or 1D numpy array which contains the column indices. The default value is None, which means all columns. :param in_cam: return points in camera frame, otherwise, return points in base frame :type rs: list or np.ndarray :type cs: list or np.ndarray :type in_cam: bool :returns: tuple (pts, colors) pts: point coordinates in world frame (shape: :math:`[N, 3]`) colors: rgb values for pts_in_cam (shape: :math:`[N, 3]`) :rtype: tuple(np.ndarray, np.ndarray) """ raise NotImplementedError def get_current_pcd(self, in_cam=True): """ Return the point cloud at current time step (one frame only) :param in_cam: return points in camera frame, otherwise, return points in base frame :type in_cam: bool :returns: tuple (pts, colors) pts: point coordinates in world frame (shape: :math:`[N, 3]`) colors: rgb values for pts_in_cam (shape: :math:`[N, 3]`) :rtype: tuple(np.ndarray, np.ndarray) """ raise NotImplementedError @property def state(self): """ Return the current pan and tilt joint angles of the robot camera. :return: pan_tilt: A list the form [pan angle, tilt angle] :rtype: list """ return self.get_state() def get_state(self): """ Return the current pan and tilt joint angles of the robot camera. :return: pan_tilt: A list the form [pan angle, tilt angle] :rtype: list """ return [self.get_pan(), self.get_tilt()] def get_pan(self): """ Return the current pan joint angle of the robot camera. :return: pan: Pan joint angle :rtype: float """ return self.pan_joint.get_joint_position() def get_tilt(self): """ Return the current tilt joint angle of the robot camera. :return: tilt: Tilt joint angle :rtype: float """ return self.tilt_joint.get_joint_position() def set_pan(self, pan, wait=True): """ Sets the pan joint angle to the specified value. :param pan: value to be set for pan joint :param wait: wait until the pan angle is set to the target angle. :type pan: float :type wait: bool """ self.pan_joint.set_joint_position(pan) # [self.sim.step() for _ in range(50)] def set_tilt(self, tilt, wait=True): """ Sets the tilt joint angle to the specified value. :param tilt: value to be set for the tilt joint :param wait: wait until the tilt angle is set to the target angle. :type tilt: float :type wait: bool """ self.tilt_joint.set_joint_position(tilt) def set_pan_tilt(self, pan, tilt, wait=True): """ Sets both the pan and tilt joint angles to the specified values. :param pan: value to be set for pan joint :param tilt: value to be set for the tilt joint :param wait: wait until the pan and tilt angles are set to the target angles. :type pan: float :type tilt: float :type wait: bool """ self.set_pan(pan) self.set_tilt(tilt) def reset(self): """ This function resets the pan and tilt joints by actuating them to their home configuration. """ self.set_pan_tilt(self.configs.CAMERA.RESET_PAN, self.configs.CAMERA.RESET_TILT)
class Scene(object): """Controls what is currently in the vrep scene. This is used for making sure that the tasks are easily reachable. This may be just replaced by environment. Responsible for moving all the objects. """ def __init__(self, pyrep: PyRep, robot: Robot, obs_config=ObservationConfig()): self._pyrep = pyrep self._robot = robot self._obs_config = obs_config self._active_task = None self._inital_task_state = None self._start_arm_joint_pos = robot.arm.get_joint_positions() self._starting_gripper_joint_pos = robot.gripper.get_joint_positions() self._workspace = Shape('workspace') self._workspace_boundary = SpawnBoundary([self._workspace]) self._cam_over_shoulder_left = VisionSensor('cam_over_shoulder_left') self._cam_over_shoulder_right = VisionSensor('cam_over_shoulder_right') self._cam_wrist = VisionSensor('cam_wrist') self._cam_over_shoulder_left_mask = VisionSensor( 'cam_over_shoulder_left_mask') self._cam_over_shoulder_right_mask = VisionSensor( 'cam_over_shoulder_right_mask') self._cam_wrist_mask = VisionSensor('cam_wrist_mask') self._has_init_task = self._has_init_episode = False self._variation_index = 0 self._initial_robot_state = (robot.arm.get_configuration_tree(), robot.gripper.get_configuration_tree()) # Set camera properties from observation config self._set_camera_properties() def load(self, task: Task) -> None: """Loads the task and positions at the centre of the workspace. :param task: The task to load in the scene. """ task.load() # Load the task in to the scene # Set at the centre of the workspace task.get_base().set_position(self._workspace.get_position()) self._inital_task_state = task.get_state() self._active_task = task self._initial_task_pose = task.boundary_root().get_orientation() self._has_init_task = self._has_init_episode = False self._variation_index = 0 def unload(self) -> None: """Clears the scene. i.e. removes all tasks. """ if self._active_task is not None: self._robot.gripper.release() if self._has_init_task: self._active_task.cleanup_() self._active_task.unload() self._active_task = None self._variation_index = 0 def init_task(self) -> None: self._active_task.init_task() self._has_init_task = True self._variation_index = 0 def init_episode(self, index: int, randomly_place: bool = True, max_attempts: int = 5) -> List[str]: """Calls the task init_episode and puts randomly in the workspace. """ self._variation_index = index if not self._has_init_task: self.init_task() # Try a few times to init and place in the workspace attempts = 0 descriptions = None while attempts < max_attempts: descriptions = self._active_task.init_episode(index) try: if (randomly_place and not self._active_task.is_static_workspace()): self._place_task() self._active_task.validate() break except (BoundaryError, WaypointError) as e: self._active_task.cleanup_() self._active_task.restore_state(self._inital_task_state) attempts += 1 if attempts >= max_attempts: raise e # Let objects come to rest [self._pyrep.step() for _ in range(STEPS_BEFORE_EPISODE_START)] self._has_init_episode = True return descriptions def reset(self) -> None: """Resets the joint angles. """ self._robot.gripper.release() arm, gripper = self._initial_robot_state self._pyrep.set_configuration_tree(arm) self._pyrep.set_configuration_tree(gripper) self._robot.arm.set_joint_positions(self._start_arm_joint_pos) self._robot.arm.set_joint_target_velocities( [0] * len(self._robot.arm.joints)) self._robot.gripper.set_joint_positions( self._starting_gripper_joint_pos) self._robot.gripper.set_joint_target_velocities( [0] * len(self._robot.gripper.joints)) if self._active_task is not None and self._has_init_task: self._active_task.cleanup_() self._active_task.restore_state(self._inital_task_state) [self._pyrep.step_ui() for _ in range(20)] self._active_task.set_initial_objects_in_scene() def get_observation(self) -> Observation: tip = self._robot.arm.get_tip() joint_forces = None if self._obs_config.joint_forces: fs = self._robot.arm.get_joint_forces() vels = self._robot.arm.get_joint_target_velocities() joint_forces = self._obs_config.joint_forces_noise.apply( np.array([-f if v < 0 else f for f, v in zip(fs, vels)])) ee_forces_flat = None if self._obs_config.gripper_touch_forces: ee_forces = self._robot.gripper.get_touch_sensor_forces() ee_forces_flat = [] for eef in ee_forces: ee_forces_flat.extend(eef) ee_forces_flat = np.array(ee_forces_flat) lsc_ob = self._obs_config.left_shoulder_camera rsc_ob = self._obs_config.right_shoulder_camera wc_ob = self._obs_config.wrist_camera obs = Observation( left_shoulder_rgb=(lsc_ob.rgb_noise.apply( self._cam_over_shoulder_left.capture_rgb()) if lsc_ob.rgb else None), left_shoulder_depth=(lsc_ob.depth_noise.apply( self._cam_over_shoulder_left.capture_depth()) if lsc_ob.depth else None), right_shoulder_rgb=(rsc_ob.rgb_noise.apply( self._cam_over_shoulder_right.capture_rgb()) if rsc_ob.rgb else None), right_shoulder_depth=(rsc_ob.depth_noise.apply( self._cam_over_shoulder_right.capture_depth()) if rsc_ob.depth else None), wrist_rgb=(wc_ob.rgb_noise.apply(self._cam_wrist.capture_rgb()) if wc_ob.rgb else None), wrist_depth=(wc_ob.depth_noise.apply( self._cam_wrist.capture_depth()) if wc_ob.depth else None), left_shoulder_mask=( self._cam_over_shoulder_left_mask.capture_rgb() if lsc_ob.mask else None), right_shoulder_mask=( self._cam_over_shoulder_right_mask.capture_rgb() if rsc_ob.mask else None), wrist_mask=(self._cam_wrist_mask.capture_rgb() if wc_ob.mask else None), joint_velocities=(self._obs_config.joint_velocities_noise.apply( np.array(self._robot.arm.get_joint_velocities())) if self._obs_config.joint_velocities else None), joint_positions=(self._obs_config.joint_positions_noise.apply( np.array(self._robot.arm.get_joint_positions())) if self._obs_config.joint_positions else None), joint_forces=(joint_forces if self._obs_config.joint_forces else None), gripper_open_amount=(( 1.0 if self._robot.gripper.get_open_amount()[0] > 0.9 else 0.0) if self._obs_config.gripper_open_amount else None), gripper_pose=(np.array(tip.get_pose()) if self._obs_config.gripper_pose else None), gripper_touch_forces=(ee_forces_flat if self._obs_config.gripper_touch_forces else None), gripper_joint_positions=( np.array(self._robot.gripper.get_joint_positions()) if self._obs_config.gripper_joint_positions else None), task_low_dim_state=(self._active_task.get_low_dim_state() if self._obs_config.task_low_dim_state else None)) obs = self._active_task.decorate_observation(obs) return obs def step(self): self._pyrep.step() self._active_task.step() def get_demo(self, record: bool = True, func=None, randomly_place: bool = True) -> List[Observation]: """Returns a demo (list of observations)""" if not self._has_init_task: self.init_task() if not self._has_init_episode: self.init_episode(self._variation_index, randomly_place=randomly_place) self._has_init_episode = False waypoints = self._active_task.get_waypoints() if len(waypoints) == 0: raise NoWaypointsError('No waypoints were found.', self._active_task) demo = [] if record: self._pyrep.step() # Need this here or get_force doesn't work... demo.append(self.get_observation()) while True: success = False for i, point in enumerate(waypoints): point.start_of_path() try: path = point.get_path() except ConfigurationPathError as e: raise DemoError( 'Could not get a path for waypoint %d.' % i, self._active_task) from e ext = point.get_ext() path.visualize() done = False success = False while not done: done = path.step() self.step() self._demo_record_step(demo, record, func) success, term = self._active_task.success() if success: break point.end_of_path() path.clear_visualization() if success: # We can quit early because we have finished the task break # TODO: need to decide how I do the gripper actions if len(ext) > 0: contains_param = False start_of_bracket = -1 gripper = self._robot.gripper if 'open_gripper(' in ext: gripper.release() start_of_bracket = ext.index('open_gripper(') + 13 contains_param = ext[start_of_bracket] != ')' if not contains_param: done = False while not done: done = gripper.actuate(1.0, 0.04) self._pyrep.step() self._active_task.step() if self._obs_config.record_gripper_closing: self._demo_record_step(demo, record, func) elif 'close_gripper(' in ext: start_of_bracket = ext.index('close_gripper(') + 14 contains_param = ext[start_of_bracket] != ')' if not contains_param: done = False while not done: done = gripper.actuate(0.0, 0.04) self._pyrep.step() self._active_task.step() if self._obs_config.record_gripper_closing: self._demo_record_step(demo, record, func) if contains_param: rest = ext[start_of_bracket:] num = float(rest[:rest.index(')')]) done = False while not done: done = gripper.actuate(num, 0.04) self._pyrep.step() self._active_task.step() if self._obs_config.record_gripper_closing: self._demo_record_step(demo, record, func) if 'close_gripper(' in ext: for g_obj in self._active_task.get_graspable_objects(): gripper.grasp(g_obj) self._demo_record_step(demo, record, func) if not self._active_task.should_repeat_waypoints() or success: break # Some tasks may need additional physics steps # (e.g. ball rowling to goal) if not success: for _ in range(10): self._pyrep.step() self._active_task.step() self._demo_record_step(demo, record, func) success, term = self._active_task.success() if success: break success, term = self._active_task.success() if not success: raise DemoError('Demo was completed, but was not successful.', self._active_task) return demo def get_observation_config(self) -> ObservationConfig: return self._obs_config def _demo_record_step(self, demo_list, record, func): if record: demo_list.append(self.get_observation()) if func is not None: func() def _set_camera_properties(self) -> None: def _set_rgb_props(rgb_cam: VisionSensor, rgb: bool, depth: bool, conf: CameraConfig): if not (rgb or depth): rgb_cam.remove() else: rgb_cam.set_resolution(conf.image_size) rgb_cam.set_render_mode(conf.render_mode) def _set_mask_props(mask_cam: VisionSensor, mask: bool, conf: CameraConfig): if not mask: mask_cam.remove() else: mask_cam.set_resolution(conf.image_size) _set_rgb_props(self._cam_over_shoulder_left, self._obs_config.left_shoulder_camera.rgb, self._obs_config.left_shoulder_camera.depth, self._obs_config.left_shoulder_camera) _set_rgb_props(self._cam_over_shoulder_right, self._obs_config.right_shoulder_camera.rgb, self._obs_config.right_shoulder_camera.depth, self._obs_config.right_shoulder_camera) _set_rgb_props(self._cam_wrist, self._obs_config.wrist_camera.rgb, self._obs_config.wrist_camera.depth, self._obs_config.wrist_camera) _set_mask_props(self._cam_over_shoulder_left_mask, self._obs_config.left_shoulder_camera.mask, self._obs_config.left_shoulder_camera) _set_mask_props(self._cam_over_shoulder_right_mask, self._obs_config.right_shoulder_camera.mask, self._obs_config.right_shoulder_camera) _set_mask_props(self._cam_wrist_mask, self._obs_config.wrist_camera.mask, self._obs_config.wrist_camera) def _place_task(self) -> None: self._workspace_boundary.clear() # Find a place in the robot workspace for task self._active_task.boundary_root().set_orientation( self._initial_task_pose) min_rot, max_rot = self._active_task.base_rotation_bounds() self._workspace_boundary.sample(self._active_task.boundary_root(), min_rotation=min_rot, max_rotation=max_rot)
def simulate(scene_dir, cls_indices): # read 3d point cloud vertices npy (for visualization) vertex_npy = np.load("mesh_data/custom_vertex.npy") # loop over all scene files in scenes directory for scene_path in os.listdir(scene_dir): # check whether it's a scene file or not if not scene_path[-3:] == 'ttt': continue # create an output directory for each scene scene_out_dir = os.path.join('./sim_data/', scene_path[:-4]) if not os.path.isdir(scene_out_dir): os.mkdir(scene_out_dir) # launch scene file SCENE_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), os.path.join(scene_dir, scene_path)) pr = PyRep() pr.launch(SCENE_FILE, headless=True) pr.start() pr.step() # define vision sensor camera = VisionSensor('cam') # set camera absolute pose to world coordinates camera.set_pose([0,0,0,0,0,0,1]) pr.step() # define scene shapes shapes = [] shapes.append(Shape('Shape1')) shapes.append(Shape('Shape2')) shapes.append(Shape('Shape3')) shapes.append(Shape('Shape4')) pr.step() print("Getting vision sensor intrinsics ...") # get vision sensor parameters cam_res = camera.get_resolution() cam_per_angle = camera.get_perspective_angle() ratio = cam_res[0]/cam_res[1] cam_angle_x = 0.0 cam_angle_y = 0.0 if (ratio > 1): cam_angle_x = cam_per_angle cam_angle_y = 2 * degrees(atan(tan(radians(cam_per_angle / 2)) / ratio)) else: cam_angle_x = 2 * degrees(atan(tan(radians(cam_per_angle / 2)) / ratio)) cam_angle_y = cam_per_angle # get vision sensor intrinsic matrix cam_focal_x = float(cam_res[0] / 2.0) / tan(radians(cam_angle_x / 2.0)) cam_focal_y = float(cam_res[1] / 2.0) / tan(radians(cam_angle_y / 2.0)) intrinsics = np.array([[cam_focal_x, 0.00000000e+00, float(cam_res[0]/2.0)], [0.00000000e+00, cam_focal_y, float(cam_res[1]/2.0)], [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]) # loop to get 5000 samples per scene for i in range(5000): print("Randomizing objects' poses ...") # set random pose to objects in the scene obj_colors = [] for shape in shapes: shape.set_pose([ random.uniform(-0.25,0.25), random.uniform(-0.25,0.25), random.uniform(0.25,2), random.uniform(-1,1), random.uniform(-1,1), random.uniform(-1,1), random.uniform(-1,1) ]) obj_colors.append([random.uniform(0,1), random.uniform(0,1), random.uniform(0,1)]) pr.step() print("Reading vision sensor RGB signal ...") # read vision sensor RGB image img = camera.capture_rgb() img = np.uint8(img * 255.0) print("Reading vision sensor depth signal ...") # read vision sensor depth image depth = camera.capture_depth() depth = np.uint8(depth * 255.0) print("Generating front mask for RGB signal ...") # generate RGB front mask front_mask = np.sum(img, axis=2) front_mask[front_mask != 0] = 255 front_mask = Image.fromarray(np.uint8(front_mask)) alpha_img = Image.fromarray(img) alpha_img.putalpha(front_mask) print("Saving sensor output ...") # save sensor output alpha_img.save(scene_out_dir+f'/{str(i).zfill(6)}-color.png') imsave(scene_out_dir+f'/{str(i).zfill(6)}-depth.png', depth) print("Getting objects' poses ...") # get poses for all objects in scene poses = [] for shape in shapes: poses.append(get_object_pose(shape, camera)) pose_mat = np.stack(poses, axis=2) # save pose visualization on RGB image img_viz = visualize_predictions(poses, cls_indices[scene_path], img, vertex_npy, intrinsics) imsave(scene_out_dir+f'/{str(i).zfill(6)}-viz.png', img_viz) print("Saving meta-data ...") # save meta-data .mat meta_dict = { 'cls_indexes' : np.array(cls_indices[scene_path]), 'intrinsic_matrix' : intrinsics, 'poses' : pose_mat } savemat(scene_out_dir+f'/{str(i).zfill(6)}-meta.mat', meta_dict) print("Performing semantic segmentation of RGB signal ...") # perform semantic segmentation of RGB image seg_img = camera.capture_rgb() seg_img = perform_segmentation(seg_img, cls_indices[scene_path], poses, vertex_npy, intrinsics) imsave(scene_out_dir+f'/{str(i).zfill(6)}-label.png', seg_img) # stop simulation pr.stop() pr.shutdown()
class Camera(): def __init__(self, name, resolution, perspective_angle, min_depth, max_depth, near_clipping=0.2, far_clipping=3): self.base = Dummy(name) self.cam_rgb = VisionSensor(name + "_rgb") self.cam_depth = VisionSensor(name + "_depth") self.cam_mask = VisionSensor(name + "_mask") self.cam_all = [self.cam_rgb, self.cam_depth, self.cam_mask] self.min_depth = min_depth self.max_depth = max_depth self.set_resolution(resolution) self.set_perspective_angle(perspective_angle) self.set_near_clipping_plane(near_clipping) self.set_far_clipping_plane(far_clipping) def get_resolution(self): return self._resolution def set_resolution(self, resolution): for cam in self.cam_all: cam.set_resolution(resolution) self._resolution = resolution def get_perspective_angle(self): return self._perspective_angle def set_perspective_angle(self, angle: float): for cam in self.cam_all: cam.set_perspective_angle(angle) self._perspective_angle = angle def set_near_clipping_plane(self, near_clipping): self.cam_rgb.set_near_clipping_plane(near_clipping) self.cam_mask.set_near_clipping_plane(near_clipping) self.cam_depth.set_near_clipping_plane(self.min_depth) def set_far_clipping_plane(self, far_clipping): self.cam_rgb.set_far_clipping_plane(far_clipping) self.cam_mask.set_far_clipping_plane(far_clipping) self.cam_depth.set_far_clipping_plane(self.max_depth) def capture_rgb(self): return np.uint8(255 * self.cam_rgb.capture_rgb())[:, :, ::-1] def capture_depth(self): return np.uint8(255 * self.cam_depth.capture_depth()) def capture_mask(self): mask = np.uint8(255 * self.cam_mask.capture_rgb()) mask = mask[:, :, 1] mask[mask != 0] = 1 return mask def set_pose(self, pose, relative_to=None): pos = pose[:3] qx, qy, qz, qw = pose[3:] quat_ori = Quaternion([qw, qx, qy, qz]) quat_rotate_z = Quaternion(axis=[0, 0, 1], angle=np.radians(180)) quat = quat_ori * quat_rotate_z qw, qx, qy, qz = quat.elements quat = [qx, qy, qz, qw] # quat = pose[3:] self.base.set_pose([*pos] + [*quat], relative_to=relative_to) def get_image(self): rgb = self.capture_rgb() depth = self.capture_depth() mask = self.capture_mask() return rgb, depth, mask