def generate_env_variation(self, num_variations=1, vary_cloth_size=True): """ Generate initial states. Note: This will also change the current states! """ max_wait_step = 500 # Maximum number of steps waiting for the cloth to stablize stable_vel_threshold = 0.1 # Cloth stable when all particles' vel are smaller than this generated_configs, generated_states = [], [] default_config = self.get_default_config() for i in range(num_variations): config = deepcopy(default_config) self.update_camera(config['camera_name'], config['camera_params'][config['camera_name']]) if vary_cloth_size: cloth_dimx, cloth_dimy = self._sample_cloth_size() config['ClothSize'] = [cloth_dimx, cloth_dimy] else: cloth_dimx, cloth_dimy = config['ClothSize'] self.set_scene(config) self.action_tool.reset([0., -1., 0.]) pickpoints = self._get_drop_point_idx( )[:2] # Pick two corners of the cloth and wait until stablize config['target_pos'] = self._get_flat_pos() self._set_to_vertical(x_low=np.random.random() * 0.2 - 0.1, height_low=np.random.random() * 0.1 + 0.1) # Get height of the cloth without the gravity. With gravity, it will be longer p1, _, p2, _ = self._get_key_point_idx() curr_pos = pyflex.get_positions().reshape(-1, 4) curr_pos[0] += np.random.random() * 0.001 # Add small jittering original_inv_mass = curr_pos[pickpoints, 3] curr_pos[ pickpoints, 3] = 0 # Set mass of the pickup point to infinity so that it generates enough force to the rest of the cloth pickpoint_pos = curr_pos[pickpoints, :3] pyflex.set_positions(curr_pos.flatten()) picker_radius = self.action_tool.picker_radius self.action_tool.update_picker_boundary([-0.3, 0.05, -0.5], [0.5, 2, 0.5]) self.action_tool.set_picker_pos(picker_pos=pickpoint_pos + np.array([0., picker_radius, 0.])) # Pick up the cloth and wait to stablize for j in range(0, max_wait_step): pyflex.step() curr_pos = pyflex.get_positions().reshape((-1, 4)) curr_vel = pyflex.get_velocities().reshape((-1, 3)) if np.alltrue(curr_vel < stable_vel_threshold) and j > 300: break curr_pos[pickpoints, :3] = pickpoint_pos pyflex.set_positions(curr_pos) curr_pos = pyflex.get_positions().reshape((-1, 4)) curr_pos[pickpoints, 3] = original_inv_mass pyflex.set_positions(curr_pos.flatten()) generated_configs.append(deepcopy(config)) print('config {}: {}'.format(i, config['camera_params'])) generated_states.append(deepcopy(self.get_state())) return generated_configs, generated_states
def random_pick_and_place(pick_num=10, pick_scale=0.01): """ Random pick a particle up and the drop it for pick_num times""" curr_pos = pyflex.get_positions().reshape(-1, 4) num_particles = curr_pos.shape[0] for i in range(pick_num): pick_id = np.random.randint(num_particles) pick_dir = np.random.random(3) * 2 - 1 pick_dir[1] = (pick_dir[1] + 1) pick_dir *= pick_scale original_inv_mass = curr_pos[pick_id, 3] for _ in range(60): curr_pos = pyflex.get_positions().reshape(-1, 4) curr_pos[pick_id, :3] += pick_dir curr_pos[pick_id, 3] = 0 pyflex.set_positions(curr_pos.flatten()) pyflex.step() # Revert mass curr_pos = pyflex.get_positions().reshape(-1, 4) curr_pos[pick_id, 3] = original_inv_mass pyflex.set_positions(curr_pos.flatten()) pyflex.step() # Wait to stabalize for _ in range(100): pyflex.step() curr_vel = pyflex.get_velocities() if np.alltrue(curr_vel < 0.01): break for _ in range(500): pyflex.step() curr_vel = pyflex.get_velocities() if np.alltrue(curr_vel < 0.01): break
def _get_obs(self): ''' return the observation based on the current flex state. ''' if self.observation_mode == 'cam_rgb': return self.get_image(self.camera_width, self.camera_height) elif self.observation_mode in ['point_cloud', 'key_point']: if self.observation_mode == 'point_cloud': particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3].flatten() pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float) pos[:len(particle_pos)] = particle_pos else: pos = np.empty(0, dtype=np.float) particle_state = pyflex.get_positions().reshape((-1, self.dim_position)) torus_center_y = np.mean(particle_state[:, 1]) torus_center_x = np.mean(particle_state[:, 0]) on_box = float(torus_center_y >= self.height) within_box_bound = float(torus_center_x > self.box_x - 0.5 * self.box_dis_x and torus_center_x < self.box_x + self.box_dis_x * 0.5) cup_state = np.array([self.box_x, self.box_dis_x, self.box_dis_z, self.height, torus_center_x, torus_center_y, on_box, within_box_bound]) return np.hstack([pos, cup_state]).flatten() else: raise NotImplementedError
def _get_obs(self): if self.observation_mode == 'cam_rgb': obs_img = self.get_image(self.camera_height, self.camera_width) goal_img = self.current_config['goal_character_img'] ret_img = np.concatenate([obs_img, goal_img], axis=2) return ret_img if self.observation_mode == 'point_cloud': particle_pos = np.array(pyflex.get_positions()).reshape( [-1, 4])[:, :3].flatten() pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float) pos[:len(particle_pos)] = particle_pos pos[len(particle_pos):] = self.current_config[ "goal_character_pos"][:, :3].flatten() elif self.observation_mode == 'key_point': particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3] keypoint_pos = particle_pos[self.key_point_indices, :3] goal_keypoint_pos = self.current_config["goal_character_pos"][ self.key_point_indices, :3] pos = np.concatenate([keypoint_pos, goal_keypoint_pos], axis=0).flatten() if self.action_mode in ['sphere', 'picker']: shapes = pyflex.get_shape_states() shapes = np.reshape(shapes, [-1, 14]) pos = np.concatenate([pos.flatten(), shapes[:, :3].flatten()]) return pos
def _get_obs(self): ''' return the observation based on the current flex state. ''' if self.observation_mode == 'cam_rgb': return self.get_image(self.camera_width, self.camera_height) elif self.observation_mode == 'point_cloud': particle_pos = np.array(pyflex.get_positions()).reshape( [-1, 4])[:, :3].flatten() pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float) pos[:len(particle_pos)] = particle_pos return pos.flatten() elif self.observation_mode == 'key_point': pos = np.empty(0, dtype=np.float) water_state = pyflex.get_positions().reshape([-1, 4]) water_num = len(water_state) in_glass = self.in_glass(water_state, self.glass_states, self.border, self.height) out_glass = water_num - in_glass in_glass = float(in_glass) / water_num out_glass = float(out_glass) / water_num cup_state = np.array([ self.glass_x, self.glass_dis_x, self.glass_dis_z, self.height, self._get_current_water_height(), in_glass, out_glass ]) return np.hstack([pos, cup_state]).flatten() else: raise NotImplementedError
def _reset(self): """ Right now only use one initial state. Need to make sure _reset always give the same result. Otherwise CEM will fail.""" if hasattr(self, 'action_tool'): particle_pos = pyflex.get_positions().reshape(-1, 4) p1, p2, p3, p4 = self._get_key_point_idx() key_point_pos = particle_pos[( p1, p2), :3] # Was changed from from p1, p4. middle_point = np.mean(key_point_pos, axis=0) self.action_tool.reset([middle_point[0], 0.1, middle_point[2]]) # picker_low = self.action_tool.picker_low # picker_high = self.action_tool.picker_high # offset_x = self.action_tool._get_pos()[0][0][0] - picker_low[0] - 0.3 # picker_low[0] += offset_x # picker_high[0] += offset_x # picker_high[0] += 1.0 # self.action_tool.update_picker_boundary(picker_low, picker_high) config = self.get_current_config() num_particles = np.prod(config['ClothSize'], dtype=int) particle_grid_idx = np.array(list(range(num_particles))).reshape( config['ClothSize'][1], config['ClothSize'][0]) # Reversed index here cloth_dimx = config['ClothSize'][0] x_split = cloth_dimx // 2 self.fold_group_a = particle_grid_idx[:, :x_split].flatten() self.fold_group_b = np.flip(particle_grid_idx, axis=1)[:, :x_split].flatten() colors = np.zeros(num_particles) colors[self.fold_group_a] = 1 # self.set_colors(colors) # TODO the phase actually changes the cloth dynamics so we do not change them for now. Maybe delete this later. pyflex.step() self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3] pos_a = self.init_pos[self.fold_group_a, :] pos_b = self.init_pos[self.fold_group_b, :] self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1)) # If cached_goal is used recover goals from flatten initial states and set one for current iteration if self.cached_goal: # Get a random state from cached goals rand_goal = np.random.randint(0, len(self.cached_goal_states)) self.goal_pos = self.cached_goal_states[rand_goal][ 'particle_pos'].reshape((-1, 4))[:, :3] #print("GOAL VS INIT", self.goal_pos.shape, self.init_pos.shape) # Compute offset to match initial position TODO # Trim goal_pos to match curr pos if self.init_pos.shape[0] < self.goal_pos.shape[0]: self.goal_pos = np.delete( self.goal_pos, np.s_[-(abs(self.goal_pos.shape[0] - self.init_pos.shape[0]) + 1):-1], 0) self.performance_init = None info = self._get_info() self.performance_init = info['performance'] return self._get_obs()
def _reset(self): """ Right now only use one initial state""" self.prev_covered_area = self._get_current_covered_area( pyflex.get_positions()) if hasattr(self, 'action_tool'): curr_pos = pyflex.get_positions() cx, cy = self._get_center_point(curr_pos) self.action_tool.reset([cx, 0.2, cy]) pyflex.step() self.init_covered_area = None info = self._get_info() self.init_covered_area = info['performance'] return self._get_obs()
def _get_info(self): goal_c_pos = self.current_config["goal_character_pos"][:, :3] current_pos = pyflex.get_positions().reshape((-1, 4))[:, :3] # way1: index matching if self.reward_type == 'index': dist = np.linalg.norm(current_pos - goal_c_pos, axis=1) reward = -np.mean(dist) if self.reward_type == 'bigraph': # way2: downsample and then use Hungarian algorithm for bipartite graph matching downsampled_cur_pos = current_pos[self.key_point_indices] downsampled_goal_pos = goal_c_pos[self.key_point_indices] W = np.zeros((len(downsampled_cur_pos), len(downsampled_cur_pos))) for idx in range(len(downsampled_cur_pos)): all_dist = np.linalg.norm(downsampled_cur_pos[idx] - downsampled_goal_pos, axis=1) W[idx, :] = all_dist row_idx, col_idx = opt.linear_sum_assignment(W) dist = W[row_idx, col_idx].sum() reward = -dist / len(downsampled_goal_pos) performance = reward performance_init = performance if self.performance_init is None else self.performance_init # Use the original performance return { 'performance': performance, 'normalized_performance': (performance - performance_init) / (self.reward_max - performance_init), }
def compute_reward(self, action=None, obs=None, set_prev_reward=False): """ The particles are splitted into two groups. The reward will be the minus average eculidean distance between each particle in group a and the crresponding particle in group b :param pos: nx4 matrix (x, y, z, inv_mass) """ # Get current pos pos = pyflex.get_positions() pos = pos.reshape((-1, 4))[:, :3] if self.cached_goal: # Trim pos to match goal cloth size if needed if self.goal_pos.shape[0] < self.init_pos.shape[0]: pos = np.delete( pos, np.s_[-(abs(pos.shape[0] - self.goal_pos.shape[0]) + 1):-1], 0) #curr_dist = (np.square(pos - self.goal_pos)).mean(axis=None) curr_dist = mse( self.goal_pos, pos) # Calc min square error between goal and curr pos else: pos_group_a = pos[self.fold_group_a] pos_group_b = pos[self.fold_group_b] pos_group_b_init = self.init_pos[self.fold_group_b] curr_dist = np.mean(np.linalg.norm(pos_group_a - pos_group_b, axis=1)) + \ 1.2 * np.mean(np.linalg.norm(pos_group_b - pos_group_b_init, axis=1)) reward = -curr_dist return reward
def _get_info(self): # Duplicate of the compute reward function! pos = pyflex.get_positions() pos = pos.reshape((-1, 4))[:, :3] pos_group_a = pos[self.fold_group_a] pos_group_b = pos[self.fold_group_b] pos_group_b_init = self.init_pos[self.fold_group_b] group_dist = np.mean(np.linalg.norm(pos_group_a - pos_group_b, axis=1)) fixation_dist = np.mean( np.linalg.norm(pos_group_b - pos_group_b_init, axis=1)) performance = -group_dist - 1.2 * fixation_dist performance_init = performance if self.performance_init is None else self.performance_init # Use the original performance info = { 'performance': performance, 'normalized_performance': (performance - performance_init) / (0. - performance_init), 'neg_group_dist': -group_dist, 'neg_fixation_dist': -fixation_dist, 'picker_pos': self.get_picker_pos() #'depth_map': self.get_depth_map() } if 'qpg' in self.action_mode: info['total_steps'] = self.action_tool.total_steps return info
def _reset(self): """ Right now only use one initial state""" if hasattr(self, 'action_tool'): self.action_tool.reset([0., 0.2, 0.]) config = self.get_current_config() self.flat_pos = self._get_flat_pos() num_particles = np.prod(config['ClothSize'], dtype=int) particle_grid_idx = np.array(list(range(num_particles))).reshape(config['ClothSize'][1], config['ClothSize'][0]) # Reversed index here cloth_dimx = config['ClothSize'][0] x_split = cloth_dimx // 2 self.fold_group_a = particle_grid_idx[:, :x_split].flatten() self.fold_group_b = np.flip(particle_grid_idx, axis=1)[:, :x_split].flatten() colors = np.zeros(num_particles) colors[self.fold_group_a] = 1 pyflex.step() self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3] pos_a = self.init_pos[self.fold_group_a, :] pos_b = self.init_pos[self.fold_group_b, :] self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1)) self.performance_init = None info = self._get_info() self.performance_init = info['performance'] return self._get_obs()
def get_depth_map(self, pix_size=0.015625): pos = pyflex.get_positions() pos = pos.reshape((-1, 4))[:, :3] dmap = np.zeros(shape=(self.camera_height, self.camera_width)) d_len_h = dmap.shape[0] d_len_w = dmap.shape[1] # Loop throught every cloth particle for i in range(pos.shape[0]): d_h = (pos[i][2] - 0.5)/pix_size # From position to pixel d_w = (pos[i][0] - 0.5)/pix_size d_h = int(d_len_h/2 + d_h) d_w = int(d_len_w/2 + d_w) # Assign correct depth to each pixel in dmap if (abs(d_h) < d_len_h) and (abs(d_w) < d_len_w): dmap[d_h][d_w] = pos[i][1] # Recover also picker height p_pos = self.get_picker_pos()[0] u = (p_pos[2] - 0.5)/pix_size # From position to pixel v = (p_pos[0] - 0.5)/pix_size u = int(d_len_h/2 + u) v = int(d_len_w/2 + v) # Assign correct depth to each pixel in dmap if (abs(u) < d_len_h) and (abs(v) < d_len_w): dmap[u][v] = p_pos[1] #img = Image.fromarray(dmap, 'I') # debug #img.show() return dmap
def _get_obs(self): if self.observation_mode == 'cam_rgb': return self.get_image(self.camera_height, self.camera_width) if self.observation_mode == 'point_cloud': particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3].flatten() pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float) pos[:len(particle_pos)] = particle_pos elif self.observation_mode == 'key_point': particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3] keypoint_pos = particle_pos[self._get_key_point_idx(), :3] pos = keypoint_pos if self.action_mode in ['sphere', 'picker', 'pickerTEO']: shapes = pyflex.get_shape_states() shapes = np.reshape(shapes, [-1, 14]) pos = np.concatenate([pos.flatten(), shapes[:, 0:3].flatten()]) return pos
def generate_alphabet_image(self): self.goal_characters_image = {} for c in self.goal_characters: default_config = self.get_default_config(c=c) goal_c_pos = self.goal_characters_position[c] self.set_scene(default_config) all_positions = pyflex.get_positions().reshape([-1, 4]) all_positions = goal_c_pos.copy() # ignore the first a few cloth particles pyflex.set_positions(all_positions) self.update_camera('default_camera', default_config['camera_params']['default_camera']) # why we need to do this? self.action_tool.reset([0., -1., 0.]) # hide picker # goal_c_img = self.get_image(self.camera_height, self.camera_width) # import time # for _ in range(50): # pyflex.step(render=True) # time.sleep(0.1) # cv2.imshow('img', self.get_image()) # cv2.waitKey() goal_c_img = self.get_image(self.camera_height, self.camera_width) # cv2.imwrite('../data/images/rope-configuration-goal-image-{}.png'.format(c), goal_c_img[:,:,::-1]) # exit() self.goal_characters_image[c] = goal_c_img.copy()
def compute_reward(self, action=None, obs=None, **kwargs): """ Reward is the matching degree to the goal character""" goal_c_pos = self.current_config["goal_character_pos"][:, :3] current_pos = pyflex.get_positions().reshape((-1, 4))[:, :3] # way1: index matching if self.reward_type == 'index': dist = np.linalg.norm(current_pos - goal_c_pos, axis=1) reward = -np.mean(dist) if self.reward_type == 'bigraph': # way2: downsample and then use Hungarian algorithm for bipartite graph matching downsampled_cur_pos = current_pos[self.key_point_indices] downsampled_goal_pos = goal_c_pos[self.key_point_indices] W = np.zeros((len(downsampled_cur_pos), len(downsampled_cur_pos))) for idx in range(len(downsampled_cur_pos)): all_dist = np.linalg.norm(downsampled_cur_pos[idx] - downsampled_goal_pos, axis=1) W[idx, :] = all_dist row_idx, col_idx = opt.linear_sum_assignment(W) dist = W[row_idx, col_idx].sum() reward = -dist / len(downsampled_goal_pos) return reward
def _get_obs(self): ''' return the observation based on the current flex state. ''' if self.observation_mode == 'cam_rgb': return self.get_image(self.camera_width, self.camera_height) elif self.observation_mode in ['point_cloud', 'key_point']: if self.observation_mode == 'point_cloud': particle_pos = np.array(pyflex.get_positions()).reshape( [-1, 4])[:, :3].flatten() pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float) pos[:len(particle_pos)] = particle_pos cup_state = np.array([ self.glass_x, self.glass_y, self.glass_rotation, self.glass_dis_x, self.glass_dis_z, self.height, self.glass_distance + self.glass_x, self.poured_height, self.poured_glass_dis_x, self.poured_glass_dis_z, self.line_box_y, self.current_config['targe_amount'] ]) else: pos = np.empty(0, dtype=np.float) cup_state = np.array([ self.glass_x, self.glass_y, self.glass_rotation, self.glass_dis_x, self.glass_dis_z, self.height, self.glass_distance + self.glass_x, self.poured_height, self.poured_glass_dis_x, self.poured_glass_dis_z, self._get_current_water_height(), self.line_box_y, self.current_config['target_amount'] ]) return np.hstack([pos, cup_state]).flatten() else: raise NotImplementedError
def reset(self, center): for i in (0, 2): offset = center[i] - (self.picker_high[i] + self.picker_low[i]) / 2. self.picker_low[i] += offset self.picker_high[i] += offset init_picker_poses = self._get_centered_picker_pos(center) for picker_pos in init_picker_poses: pyflex.add_sphere(self.picker_radius, picker_pos, [1, 0, 0, 0]) pos = pyflex.get_shape_states( ) # Need to call this to update the shape collision pyflex.set_shape_states(pos) self.picked_particles = [None] * self.num_picker shape_state = np.array(pyflex.get_shape_states()).reshape(-1, 14) centered_picker_pos = self._get_centered_picker_pos(center) for (i, centered_picker_pos) in enumerate(centered_picker_pos): shape_state[i] = np.hstack([ centered_picker_pos, centered_picker_pos, [1, 0, 0, 0], [1, 0, 0, 0] ]) pyflex.set_shape_states(shape_state) # pyflex.step() # Remove this as having an additional step here may affect the cloth drop env self.particle_inv_mass = pyflex.get_positions().reshape(-1, 4)[:, 3]
def get_state(self): pos = pyflex.get_positions() vel = pyflex.get_velocities() shape_pos = pyflex.get_shape_states() phase = pyflex.get_phases() camera_params = copy.deepcopy(self.camera_params) return {'particle_pos': pos, 'particle_vel': vel, 'shape_pos': shape_pos, 'phase': phase, 'camera_params': camera_params, 'config_id': self.current_config_id}
def rotate_particles(self, angle): pos = pyflex.get_positions().reshape(-1, 4) center = np.mean(pos, axis=0) pos -= center new_pos = pos.copy() new_pos[:, 0] = (np.cos(angle) * pos[:, 0] - np.sin(angle) * pos[:, 2]) new_pos[:, 2] = (np.sin(angle) * pos[:, 0] + np.cos(angle) * pos[:, 2]) new_pos += center pyflex.set_positions(new_pos)
def _set_to_vertical(self, x_low, height_low): curr_pos = pyflex.get_positions().reshape((-1, 4)) vertical_pos = self._get_vertical_pos(x_low, height_low) curr_pos[:, :3] = vertical_pos max_height = np.max(curr_pos[:, 1]) if max_height < 0.5: curr_pos[:, 1] += 0.5 - max_height pyflex.set_positions(curr_pos) pyflex.step()
def center_object(): """ Center the object to be at the origin NOTE: call a pyflex.set_positions and then pyflex.step """ pos = pyflex.get_positions().reshape(-1, 4) pos[:, [0, 2]] -= np.mean(pos[:, [0, 2]], axis=0, keepdims=True) pyflex.set_positions(pos.flatten()) pyflex.step()
def get_state(self): ''' get the postion, velocity of flex particles, and postions of flex shapes. ''' particle_pos = pyflex.get_positions() particle_vel = pyflex.get_velocities() shape_position = pyflex.get_shape_states() return {'particle_pos': particle_pos, 'particle_vel': particle_vel, 'shape_pos': shape_position, 'box_x': self.box_x, 'box_states': self.box_states, 'box_params': self.box_params, 'config_id': self.current_config_id}
def _reset(self): """ Right now only use one initial state. Need to make sure _reset always give the same result. Otherwise CEM will fail.""" if hasattr(self, 'action_tool'): particle_pos = pyflex.get_positions().reshape(-1, 4) p1, p2, p3, p4 = self._get_key_point_idx() key_point_pos = particle_pos[( p1, p2), :3] # Was changed from from p1, p4. middle_point = np.mean(key_point_pos, axis=0) self.action_tool.reset([middle_point[0], 0.1, middle_point[2]]) # picker_low = self.action_tool.picker_low # picker_high = self.action_tool.picker_high # offset_x = self.action_tool._get_pos()[0][0][0] - picker_low[0] - 0.3 # picker_low[0] += offset_x # picker_high[0] += offset_x # picker_high[0] += 1.0 # self.action_tool.update_picker_boundary(picker_low, picker_high) config = self.get_current_config() num_particles = np.prod(config['ClothSize'], dtype=int) particle_grid_idx = np.array(list(range(num_particles))).reshape( config['ClothSize'][1], config['ClothSize'][0]) # Reversed index here cloth_dimx = config['ClothSize'][0] x_split = cloth_dimx // 2 self.fold_group_a = particle_grid_idx[:, :x_split].flatten() self.fold_group_b = np.flip(particle_grid_idx, axis=1)[:, :x_split].flatten() colors = np.zeros(num_particles) colors[self.fold_group_a] = 1 # self.set_colors(colors) # TODO the phase actually changes the cloth dynamics so we do not change them for now. Maybe delete this later. pyflex.step() self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3] pos_a = self.init_pos[self.fold_group_a, :] pos_b = self.init_pos[self.fold_group_b, :] self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1)) self.performance_init = None info = self._get_info() self.performance_init = info['performance'] return self._get_obs()
def _reset(self): """ Right now only use one initial state""" self.prev_dist = self._get_current_dist(pyflex.get_positions()) if hasattr(self, 'action_tool'): particle_pos = pyflex.get_positions().reshape(-1, 4) drop_point_pos = particle_pos[self._get_drop_point_idx(), :3] middle_point = np.mean(drop_point_pos, axis=0) self.action_tool.reset( middle_point) # middle point is not really useful picker_radius = self.action_tool.picker_radius self.action_tool.update_picker_boundary([-0.3, 0.5, -0.5], [0.5, 2, 0.5]) self.action_tool.set_picker_pos(picker_pos=drop_point_pos + np.array([0., picker_radius, 0.])) # self.action_tool.visualize_picker_boundary() self.performance_init = None info = self._get_info() self.performance_init = info['performance'] return self._get_obs()
def _get_info(self): particle_pos = pyflex.get_positions() curr_dist = self._get_current_dist(particle_pos) performance = -curr_dist performance_init = performance if self.performance_init is None else self.performance_init # Use the original performance return { 'performance': performance, 'normalized_performance': (performance - performance_init) / (0. - performance_init) }
def _get_obs(self): if self.observation_mode == 'cam_rgb': return self.get_image(self.camera_height, self.camera_width) if self.observation_mode == 'point_cloud': particle_pos = np.array(pyflex.get_positions()).reshape( [-1, 4])[:, :3].flatten() pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float) pos[:len(particle_pos)] = particle_pos elif self.observation_mode == 'key_point': particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3] keypoint_pos = particle_pos[self.key_point_indices, :3] pos = keypoint_pos.flatten() # more_info = np.array([self.rope_length, self._get_endpoint_distance()]) # pos = np.hstack([more_info, pos]) # print("in _get_obs, pos is: ", pos) if self.action_mode in ['sphere', 'picker']: shapes = pyflex.get_shape_states() shapes = np.reshape(shapes, [-1, 14]) pos = np.concatenate([pos.flatten(), shapes[:, 0:3].flatten()]) return pos
def _reset(self): """ Right now only use one initial state""" if hasattr(self, 'action_tool'): particle_pos = pyflex.get_positions().reshape(-1, 4) drop_point_pos = particle_pos[self._get_drop_point_idx(), :3] middle_point = np.mean(drop_point_pos, axis=0) self.action_tool.reset( middle_point) # middle point is not really useful picker_radius = self.action_tool.picker_radius self.action_tool.update_picker_boundary([-0.3, 0.5, -0.5], [0.5, 2, 0.5]) self.action_tool.set_picker_pos(picker_pos=drop_point_pos + np.array([0., picker_radius, 0.])) config = self.get_current_config() num_particles = np.prod(config['ClothSize'], dtype=int) particle_grid_idx = np.array(list(range(num_particles))).reshape( config['ClothSize'][1], config['ClothSize'][0]) # Reversed index here cloth_dimx = config['ClothSize'][0] x_split = cloth_dimx // 2 self.fold_group_a = particle_grid_idx[:, :x_split].flatten() self.fold_group_b = np.flip(particle_grid_idx, axis=1)[:, :x_split].flatten() colors = np.zeros(num_particles) colors[self.fold_group_a] = 1 pyflex.step() self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3] pos_a = self.init_pos[self.fold_group_a, :] pos_b = self.init_pos[self.fold_group_b, :] self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1)) self.performance_init = None info = self._get_info() self.performance_init = info['performance'] return self._get_obs()
def compute_reward(self, action=None, obs=None, set_prev_reward=False): """ The particles are splitted into two groups. The reward will be the minus average eculidean distance between each particle in group a and the crresponding particle in group b :param pos: nx4 matrix (x, y, z, inv_mass) """ pos = pyflex.get_positions() pos = pos.reshape((-1, 4))[:, :3] pos_group_a = pos[self.fold_group_a] pos_group_b = pos[self.fold_group_b] pos_group_b_init = self.flat_pos[self.fold_group_b] curr_dist = np.mean(np.linalg.norm(pos_group_a - pos_group_b, axis=1)) + 1.2 * np.linalg.norm(np.mean(pos_group_b - pos_group_b_init, axis=1)) reward = -curr_dist return reward
def _get_info(self): # Duplicate of the compute reward function! particle_pos = pyflex.get_positions() curr_covered_area = self._get_current_covered_area(particle_pos) init_covered_area = curr_covered_area if self.init_covered_area is None else self.init_covered_area max_covered_area = self.get_current_config()['flatten_area'] info = { 'performance': curr_covered_area, 'normalized_performance': (curr_covered_area - init_covered_area) / (max_covered_area - init_covered_area), } if 'qpg' in self.action_mode: info['total_steps'] = self.action_tool.total_steps return info
def generate_env_variation(self, num_variations=2, vary_cloth_size=True): """ Generate initial states. Note: This will also change the current states! """ max_wait_step = 1000 # Maximum number of steps waiting for the cloth to stablize stable_vel_threshold = 0.2 # Cloth stable when all particles' vel are smaller than this generated_configs, generated_states = [], [] default_config = self.get_default_config() default_config['flip_mesh'] = 1 for i in range(num_variations): config = deepcopy(default_config) self.update_camera(config['camera_name'], config['camera_params'][config['camera_name']]) if vary_cloth_size: cloth_dimx, cloth_dimy = self._sample_cloth_size() config['ClothSize'] = [cloth_dimx, cloth_dimy] else: cloth_dimx, cloth_dimy = config['ClothSize'] self.set_scene(config) self.action_tool.reset([0., -1., 0.]) pos = pyflex.get_positions().reshape(-1, 4) pos[:, :3] -= np.mean(pos, axis=0)[:3] if self.action_mode in ['sawyer', 'franka' ]: # Take care of the table in robot case pos[:, 1] = 0.57 else: pos[:, 1] = 0.005 pos[:, 3] = 1 pyflex.set_positions(pos.flatten()) pyflex.set_velocities(np.zeros_like(pos)) for _ in range(5): # In case if the cloth starts in the air pyflex.step() for wait_i in range(max_wait_step): pyflex.step() curr_vel = pyflex.get_velocities() if np.alltrue(np.abs(curr_vel) < stable_vel_threshold): break center_object() angle = (np.random.random() - 0.5) * np.pi / 2 self.rotate_particles(angle) generated_configs.append(deepcopy(config)) print('config {}: {}'.format(i, config['camera_params'])) generated_states.append(deepcopy(self.get_state())) return generated_configs, generated_states