def _step(self, action): ''' action: np.ndarray of dim 1x3, (x, y, theta). (x, y) specifies the floor center coordinate, and theta specifies the rotation. ''' # make action as increasement, clip its range move = action[:2] rotate = action[2] move = np.clip(move, a_min=self.action_space.low[0], a_max=self.action_space.high[0]) rotate = np.clip(rotate, a_min=self.action_space.low[2], a_max=self.action_space.high[2]) dx, dy, dtheta = move[0], move[1], rotate x, y, theta = self.glass_x + dx, self.glass_y + dy, self.glass_rotation + dtheta # check if the movement of the pouring glass collide with the poured glass. # the action only take effects if there is no collision new_states = self.rotate_glass(self.glass_states, x, y, theta) if not self.judge_glass_collide( new_states, theta) and self.above_floor(new_states, theta): self.glass_states = new_states self.glass_x, self.glass_y, self.glass_rotation = x, y, theta else: # invalid move, old state becomes the same as the current state self.glass_states[:, 3:6] = self.glass_states[:, :3].copy() self.glass_states[:, 10:] = self.glass_states[:, 6:10].copy() # pyflex takes a step to update the glass and the water fluid self.set_shape_states(self.glass_states, self.poured_glass_states) pyflex.step(render=True) self.inner_step += 1
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 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 _step(self, action): if self.action_mode.startswith('picker'): self.action_tool.step(action) pyflex.step() else: raise NotImplementedError return
def _step(self, action): self.action_tool.step(action) if self.action_mode in ['sawyer', 'franka']: pyflex.step(self.action_tool.next_action) else: pyflex.step() return
def step(self, action): """ action: Array of pick_num x 4. For each picker, the action should be [x, y, z, pick/drop]. The picker will then first pick/drop, and keep the pick/drop state while moving towards x, y, x. """ total_steps = 0 action = action.reshape(-1, 4) curr_pos = np.array(pyflex.get_shape_states()).reshape(-1, 14)[:, :3] end_pos = np.vstack([ self._apply_picker_boundary(picker_pos) for picker_pos in action[:, :3] ]) dist = np.linalg.norm(curr_pos - end_pos, axis=1) num_step = np.max(np.ceil(dist / self.delta_move)) if num_step < 0.1: return delta = (end_pos - curr_pos) / num_step norm_delta = np.linalg.norm(delta) for i in range( int(min(num_step, 300)) ): # The maximum number of steps allowed for one pick and place curr_pos = np.array(pyflex.get_shape_states()).reshape(-1, 14)[:, :3] dist = np.linalg.norm(end_pos - curr_pos, axis=1) if np.alltrue(dist < norm_delta): delta = end_pos - curr_pos super().step(np.hstack([delta, action[:, 3].reshape(-1, 1)])) pyflex.step() total_steps += 1 if self.env is not None and self.env.recording: self.env.video_frames.append(self.env.render(mode='rgb_array')) if np.alltrue(dist < self.delta_move): break return total_steps
def _reset(self): super()._reset() self.performance_init = None info = self._get_info() self.performance_init = info['performance'] pyflex.step() return self._get_obs()
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 _step(self, action): #if self.action_mode == 'pickerTEO': # action = self.action_tool.step_teo(action) self.action_tool.step(action) if self.action_mode in ['sawyer', 'franka']: print(self.action_tool.next_action) pyflex.step(self.action_tool.next_action) else: pyflex.step()
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 _reset(self): ''' reset to environment to the initial state. return the initial observation. ''' self.inner_step = 0 self.performance_init = None info = self._get_info() self.performance_init = info['performance'] pyflex.step(render=True) return self._get_obs()
def set_state(self, state_dic): ''' set the postion, velocity of flex particles, and postions of flex shapes. ''' self.box_params = state_dic['box_params'] pyflex.set_positions(state_dic["particle_pos"]) pyflex.set_velocities(state_dic["particle_vel"]) pyflex.set_shape_states(state_dic["shape_pos"]) self.box_x = state_dic['box_x'] self.box_states = state_dic['box_states'] for _ in range(5): pyflex.step()
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 set_state(self, state_dic): ''' set the postion, velocity of flex particles, and postions of flex shapes. ''' pyflex.set_positions(state_dic["particle_pos"]) pyflex.set_velocities(state_dic["particle_vel"]) pyflex.set_shape_states(state_dic["shape_pos"]) self.glass_x = state_dic['glass_x'] self.glass_y = state_dic['glass_y'] self.glass_rotation = state_dic['glass_rotation'] self.glass_states = state_dic['glass_states'] self.poured_glass_states = state_dic['poured_glass_states'] for _ in range(5): pyflex.step()
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 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
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 _step(self, action): ''' action: np.ndarray of dim 1x1, dx, which specifies how much to move on the x-axis. ''' # make action as increasement, clip its range dx = action[0] dx = np.clip(dx, a_min=self.action_space.low[0], a_max=self.action_space.high[0]) x = self.box_x + dx # move the box new_states = self.move_box(self.box_states, x) self.box_states = new_states self.box_x = x self.box_x = np.clip(self.box_x, a_min=self.min_x, a_max=self.max_x) # pyflex takes a step to update the box and the water fluid pyflex.set_shape_states(self.box_states) pyflex.step() self.inner_step += 1
def set_shape_states(self, glass_states, poured_glass_states): all_states = np.concatenate((glass_states, poured_glass_states), axis=0) if self.line_box_x is not None: quat = quatFromAxisAngle([0, 0, -1.], 0.) indicator_box_line_states = np.zeros((1, self.dim_shape_state)) indicator_box_line_states[0, :3] = np.array( [self.line_box_x, self.line_box_y, 0.]) indicator_box_line_states[0, 3:6] = np.array( [self.line_box_x, self.line_box_y, 0.]) indicator_box_line_states[:, 6:10] = quat indicator_box_line_states[:, 10:] = quat all_states = np.concatenate( (all_states, indicator_box_line_states), axis=0) pyflex.set_shape_states(all_states) if self.line_box_x is not None: pyflex.step(render=True)
def set_scene(self, config, states=None): ''' Construct the passing water scence. ''' # create fluid super().set_scene(config) # do not sample fluid parameters, as it's very likely to generate very strange fluid # compute box params if states is None: self.set_box_params(config) else: box_params = states['box_params'] self.height = box_params['height'] self.box_dis_x = box_params['box_dis_x'] self.box_dis_z = box_params['box_dis_z'] self.box_params = box_params self.x_center = 0 # create box self.create_box(self.box_dis_x, self.box_dis_z, self.height) # move box to be at ground or on the table self.box_states = self.init_box_state(self.x_center, 0, self.box_dis_x, self.box_dis_z, self.height) pyflex.set_shape_states(self.box_states) # record box floor center x self.box_x = self.x_center # no cached init states passed in if states is None: for _ in range(50): pyflex.step() pyflex.render() return True else: # set to passed-in cached init states self.set_state(states)
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 _set_to_folded(self): 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 fold_group_a = particle_grid_idx[:, :x_split].flatten() fold_group_b = np.flip(particle_grid_idx, axis=1)[:, :x_split].flatten() curr_pos = pyflex.get_positions().reshape((-1, 4)) curr_pos[fold_group_a, :] = curr_pos[fold_group_b, :].copy() curr_pos[ fold_group_a, 1] += 0.05 # group a particle position made tcurr_pos[self.fold_group_b, 1] + 0.05e at top of group b position. pyflex.set_positions(curr_pos) for i in range(10): pyflex.step() return self._get_info()['performance']
def step(self, action): """ Action is in 5D: (u,v) the start of the pick in image coordinate; (dx, dy, dz): the relative position of the place w.r.t. the pick""" u, v = action[:2] u = ((u + 1.) * 0.5) * self.image_size[0] v = ((v + 1.) * 0.5) * self.image_size[1] x, y, z = self._get_world_coor_from_image(u, v) y += 0.01 dx, dy, dz = action[2:] st_high = np.array([x, 0.2, z, 0]) st = np.array([x, y, z, 0]) en = st + np.array([dx, dy, dz, 1]) # print('st:', st) if self.full: self.total_steps += super().step(st_high) self.total_steps += super().step(st) self.total_steps += super().step(en) en[3] = 0 # Drop cloth # Unpick all particles _, particle_pos = self._get_pos() new_particle_pos = particle_pos.copy() for i in range(self.num_picker): if self.picked_particles[i] is not None: new_particle_pos[ self.picked_particles[i], 3] = self.particle_inv_mass[ self.picked_particles[i]] # Revert the mass self.picked_particles[i] = None pyflex.set_positions(new_particle_pos) for i in range(20): pyflex.step() if self.env is not None and self.env.recording: self.env.video_frames.append( self.env.render(mode='rgb_array')) self.total_steps += 20 else: raise NotImplementedError return self.total_steps
def set_scene(self, config, states=None, create_only=False): ''' Construct the pouring water scence. ''' # create fluid super().set_scene( config ) # do not sample fluid parameters, as it's very likely to generate very strange fluid # compute glass params if states is None: self.set_pouring_glass_params(config["glass"]) self.set_poured_glass_params(config["glass"]) else: glass_params = states['glass_params'] self.border = glass_params['border'] self.height = glass_params['height'] self.glass_dis_x = glass_params['glass_dis_x'] self.glass_dis_z = glass_params['glass_dis_z'] self.glass_distance = glass_params['glass_distance'] self.poured_border = glass_params['poured_border'] self.poured_height = glass_params['poured_height'] self.poured_glass_dis_x = glass_params['poured_glass_dis_x'] self.poured_glass_dis_z = glass_params['poured_glass_dis_z'] self.glass_params = glass_params # create pouring glass & poured glass self.create_glass(self.glass_dis_x, self.glass_dis_z, self.height, self.border) self.create_glass(self.poured_glass_dis_x, self.poured_glass_dis_z, self.poured_height, self.poured_border) # move pouring glass to be at ground self.glass_states = self.init_glass_state(self.x_center, 0, self.glass_dis_x, self.glass_dis_z, self.height, self.border) # move poured glass to be at ground self.poured_glass_states = self.init_glass_state( self.x_center + self.glass_distance, 0, self.poured_glass_dis_x, self.poured_glass_dis_z, self.poured_height, self.poured_border) self.set_shape_states(self.glass_states, self.poured_glass_states) # record glass floor center x, y, and rotation self.glass_x = self.x_center if self.action_mode == 'rotation_bottom': self.glass_y = 0 elif self.action_mode == 'rotation_top': self.glass_y = 0.5 * self.border + self.height self.glass_rotation = 0 # only create the glass and water, without setting their states # this is only used in the pourwater amount env. if create_only: return # no cached init states passed in if states is None: fluid_pos = np.ones((self.particle_num, self.dim_position)) # move water all inside the glass fluid_radius = self.fluid_params['radius'] * self.fluid_params[ 'rest_dis_coef'] fluid_dis = np.array( [1.0 * fluid_radius, fluid_radius * 0.5, 1.0 * fluid_radius]) lower_x = self.glass_params['glass_x_center'] - self.glass_params[ 'glass_dis_x'] / 2. + self.glass_params['border'] lower_z = -self.glass_params[ 'glass_dis_z'] / 2 + self.glass_params['border'] lower_y = self.glass_params['border'] if self.action_mode in ['sawyer', 'franka']: lower_y += 0.56 # NOTE: robotics table lower = np.array([lower_x, lower_y, lower_z]) cnt = 0 rx = int(self.fluid_params['dim_x'] * 1) ry = int(self.fluid_params['dim_y'] * 1) rz = int(self.fluid_params['dim_z'] / 1) for x in range(rx): for y in range(ry): for z in range(rz): fluid_pos[cnt][:3] = lower + np.array( [x, y, z]) * fluid_dis # + np.random.rand() * 0.01 cnt += 1 pyflex.set_positions(fluid_pos) print("stablize water!") for _ in range(100): pyflex.step() state_dic = self.get_state() water_state = state_dic['particle_pos'].reshape( (-1, self.dim_position)) in_glass = self.in_glass(water_state, self.glass_states, self.border, self.height) not_in_glass = 1 - in_glass not_total_num = np.sum(not_in_glass) while not_total_num > 0: max_height_now = np.max(water_state[:, 1]) fluid_dis = np.array( [1.0 * fluid_radius, fluid_radius * 1, 1.0 * fluid_radius]) lower_x = self.glass_params[ 'glass_x_center'] - self.glass_params['glass_dis_x'] / 4 lower_z = -self.glass_params['glass_dis_z'] / 4 lower_y = max_height_now lower = np.array([lower_x, lower_y, lower_z]) cnt = 0 dim_x = config['fluid']['dim_x'] dim_z = config['fluid']['dim_z'] for w_idx in range(len(water_state)): if not in_glass[w_idx]: water_state[w_idx][:3] = lower + fluid_dis * np.array([ cnt % dim_x, cnt // (dim_x * dim_z), (cnt // dim_x) % dim_z ]) cnt += 1 pyflex.set_positions(water_state) for _ in range(40): pyflex.step() state_dic = self.get_state() water_state = state_dic['particle_pos'].reshape( (-1, self.dim_position)) in_glass = self.in_glass(water_state, self.glass_states, self.border, self.height) not_in_glass = 1 - in_glass not_total_num = np.sum(not_in_glass) for _ in range(30): pyflex.step() else: # set to passed-in cached init states self.set_state(states)
def _get_center_point(self, pos): pos = np.reshape(pos, [-1, 4]) min_x = np.min(pos[:, 0]) min_y = np.min(pos[:, 2]) max_x = np.max(pos[:, 0]) max_y = np.max(pos[:, 2]) return 0.5 * (min_x + max_x), 0.5 * (min_y + max_y) if __name__ == '__main__': env = RopeNewEnv(observation_mode='key_point', action_mode='picker', num_picker=2, render=True, headless=False, horizon=75, action_repeat=8, num_variations=10, use_cached_states=False, save_cached_states=False, deterministic=False) env.reset(config=env.get_default_config()) for i in range(1000): print(i) print("right before pyflex step") pyflex.step() print("right after pyflex step") print("right before pyflex render") pyflex.render() print("right after pyflex render")
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 = 300 # Maximum number of steps waiting for the cloth to stablize stable_vel_threshold = 0.01 # 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.]) 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)) pyflex.step() num_particle = cloth_dimx * cloth_dimy pickpoint = random.randint(0, num_particle - 1) curr_pos = pyflex.get_positions() original_inv_mass = curr_pos[pickpoint * 4 + 3] curr_pos[ pickpoint * 4 + 3] = 0 # Set the mass of the pickup point to infinity so that it generates enough force to the rest of the cloth pickpoint_pos = curr_pos[pickpoint * 4:pickpoint * 4 + 3].copy( ) # Pos of the pickup point is fixed to this point pickpoint_pos[1] += np.random.random(1) * 0.5 + 0.5 pyflex.set_positions(curr_pos) # Pick up the cloth and wait to stablize for j in range(0, max_wait_step): curr_pos = pyflex.get_positions() curr_vel = pyflex.get_velocities() curr_pos[pickpoint * 4:pickpoint * 4 + 3] = pickpoint_pos curr_vel[pickpoint * 3:pickpoint * 3 + 3] = [0, 0, 0] pyflex.set_positions(curr_pos) pyflex.set_velocities(curr_vel) pyflex.step() if np.alltrue( np.abs(curr_vel) < stable_vel_threshold) and j > 5: break # Drop the cloth and wait to stablize curr_pos = pyflex.get_positions() curr_pos[pickpoint * 4 + 3] = original_inv_mass pyflex.set_positions(curr_pos) for _ in range(max_wait_step): pyflex.step() curr_vel = pyflex.get_velocities() if np.alltrue(curr_vel < stable_vel_threshold): break center_object() if self.action_mode == 'sphere' or self.action_mode.startswith( 'picker'): curr_pos = pyflex.get_positions() self.action_tool.reset(curr_pos[pickpoint * 4:pickpoint * 4 + 3] + [0., 0.2, 0.]) generated_configs.append(deepcopy(config)) generated_states.append(deepcopy(self.get_state())) self.current_config = config # Needed in _set_to_flatten function generated_configs[-1]['flatten_area'] = self._set_to_flatten( ) # Record the maximum flatten area print('config {}: camera params {}, flatten area: {}'.format( i, config['camera_params'], generated_configs[-1]['flatten_area'])) return generated_configs, generated_states
def gen_PyFleX(info): env, root_num = info['env'], info['root_num'] thread_idx, data_dir, data_names = info['thread_idx'], info['data_dir'], info['data_names'] n_rollout, n_instance = info['n_rollout'], info['n_instance'] time_step, time_step_clip = info['time_step'], info['time_step_clip'] shape_state_dim, dt = info['shape_state_dim'], info['dt'] env_idx = info['env_idx'] np.random.seed(round(time.time() * 1000 + thread_idx) % 2**32) ### NOTE: we might want to fix the seed for reproduction # positions, velocities if env_idx == 5: # RiceGrip stats = [init_stat(6), init_stat(6)] else: stats = [init_stat(3), init_stat(3)] import pyflex pyflex.init() for i in range(n_rollout): if i % 10 == 0: print("%d / %d" % (i, n_rollout)) rollout_idx = thread_idx * n_rollout + i rollout_dir = os.path.join(data_dir, str(rollout_idx)) os.system('mkdir -p ' + rollout_dir) if env == 'FluidFall': scene_params = np.zeros(1) pyflex.set_scene(env_idx, scene_params, thread_idx) n_particles = pyflex.get_n_particles() positions = np.zeros((time_step, n_particles, 3), dtype=np.float32) velocities = np.zeros((time_step, n_particles, 3), dtype=np.float32) for j in range(time_step_clip): p_clip = pyflex.get_positions().reshape(-1, 4)[:, :3] pyflex.step() for j in range(time_step): positions[j] = pyflex.get_positions().reshape(-1, 4)[:, :3] if j == 0: velocities[j] = (positions[j] - p_clip) / dt else: velocities[j] = (positions[j] - positions[j - 1]) / dt pyflex.step() data = [positions[j], velocities[j]] store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5')) elif env == 'BoxBath': # BoxBath scene_params = np.zeros(1) pyflex.set_scene(env_idx, scene_params, thread_idx) n_particles = pyflex.get_n_particles() positions = np.zeros((time_step, n_particles, 3), dtype=np.float32) velocities = np.zeros((time_step, n_particles, 3), dtype=np.float32) for j in range(time_step_clip): pyflex.step() p = pyflex.get_positions().reshape(-1, 4)[:64, :3] clusters = [] st_time = time.time() kmeans = MiniBatchKMeans(n_clusters=root_num[0][0], random_state=0).fit(p) # print('Time on kmeans', time.time() - st_time) clusters.append([[kmeans.labels_]]) # centers = kmeans.cluster_centers_ ref_rigid = p for j in range(time_step): positions[j] = pyflex.get_positions().reshape(-1, 4)[:, :3] # apply rigid projection to ground truth XX = ref_rigid YY = positions[j, :64] # print("MSE init", np.mean(np.square(XX - YY))) X = XX.copy().T Y = YY.copy().T mean_X = np.mean(X, 1, keepdims=True) mean_Y = np.mean(Y, 1, keepdims=True) X = X - mean_X Y = Y - mean_Y C = np.dot(X, Y.T) U, S, Vt = np.linalg.svd(C) D = np.eye(3) D[2, 2] = np.linalg.det(np.dot(Vt.T, U.T)) R = np.dot(Vt.T, np.dot(D, U.T)) t = mean_Y - np.dot(R, mean_X) YY_fitted = (np.dot(R, XX.T) + t).T # print("MSE fit", np.mean(np.square(YY_fitted - YY))) positions[j, :64] = YY_fitted if j > 0: velocities[j] = (positions[j] - positions[j - 1]) / dt pyflex.step() data = [positions[j], velocities[j], clusters] store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5')) elif env == 'FluidShake': # if env is FluidShake height = 1.0 border = 0.025 dim_x = rand_int(10, 12) dim_y = rand_int(15, 20) dim_z = 3 x_center = rand_float(-0.2, 0.2) x = x_center - (dim_x-1)/2.*0.055 y = 0.055/2. + border + 0.01 z = 0. - (dim_z-1)/2.*0.055 box_dis_x = dim_x * 0.055 + rand_float(0., 0.3) box_dis_z = 0.2 scene_params = np.array([x, y, z, dim_x, dim_y, dim_z, box_dis_x, box_dis_z]) pyflex.set_scene(env_idx, scene_params, 0) boxes = calc_box_init_FluidShake(box_dis_x, box_dis_z, height, border) for i in range(len(boxes)): halfEdge = boxes[i][0] center = boxes[i][1] quat = boxes[i][2] pyflex.add_box(halfEdge, center, quat) n_particles = pyflex.get_n_particles() n_shapes = pyflex.get_n_shapes() # print("n_particles", n_particles) # print("n_shapes", n_shapes) positions = np.zeros((time_step, n_particles + n_shapes, 3), dtype=np.float32) velocities = np.zeros((time_step, n_particles + n_shapes, 3), dtype=np.float32) shape_quats = np.zeros((time_step, n_shapes, 4), dtype=np.float32) x_box = x_center v_box = 0. for j in range(time_step_clip): x_box_last = x_box x_box += v_box * dt shape_states_ = calc_shape_states_FluidShake( x_box, x_box_last, scene_params[-2:], height, border) pyflex.set_shape_states(shape_states_) pyflex.step() for j in range(time_step): x_box_last = x_box x_box += v_box * dt v_box += rand_float(-0.15, 0.15) - x_box * 0.1 shape_states_ = calc_shape_states_FluidShake( x_box, x_box_last, scene_params[-2:], height, border) pyflex.set_shape_states(shape_states_) positions[j, :n_particles] = pyflex.get_positions().reshape(-1, 4)[:, :3] shape_states = pyflex.get_shape_states().reshape(-1, shape_state_dim) for k in range(n_shapes): positions[j, n_particles + k] = shape_states[k, :3] shape_quats[j, k] = shape_states[k, 6:10] if j > 0: velocities[j] = (positions[j] - positions[j - 1]) / dt pyflex.step() # NOTE: 1) particle + glass wall positions, 2) particle + glass wall velocitys, 3) glass wall rotations, 4) scenen parameters data = [positions[j], velocities[j], shape_quats[j], scene_params] store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5')) elif env == 'RiceGrip': # if env is RiceGrip # repeat the grip for R times R = 3 gripper_config = sample_control_RiceGrip() if i % R == 0: ### set scene # x, y, z: [8.0, 10.0] # clusterStiffness: [0.3, 0.7] # clusterPlasticThreshold: [0.00001, 0.0005] # clusterPlasticCreep: [0.1, 0.3] x = rand_float(8.0, 10.0) y = rand_float(8.0, 10.0) z = rand_float(8.0, 10.0) clusterStiffness = rand_float(0.3, 0.7) clusterPlasticThreshold = rand_float(0.00001, 0.0005) clusterPlasticCreep = rand_float(0.1, 0.3) scene_params = np.array([x, y, z, clusterStiffness, clusterPlasticThreshold, clusterPlasticCreep]) pyflex.set_scene(env_idx, scene_params, thread_idx) scene_params[4] *= 1000. halfEdge = np.array([0.15, 0.8, 0.15]) center = np.array([0., 0., 0.]) quat = np.array([1., 0., 0., 0.]) pyflex.add_box(halfEdge, center, quat) pyflex.add_box(halfEdge, center, quat) n_particles = pyflex.get_n_particles() n_shapes = pyflex.get_n_shapes() positions = np.zeros((time_step, n_particles + n_shapes, 6), dtype=np.float32) velocities = np.zeros((time_step, n_particles + n_shapes, 6), dtype=np.float32) shape_quats = np.zeros((time_step, n_shapes, 4), dtype=np.float32) for j in range(time_step_clip): shape_states = calc_shape_states_RiceGrip(0, dt, shape_state_dim, gripper_config) pyflex.set_shape_states(shape_states) pyflex.step() p = pyflex.get_positions().reshape(-1, 4)[:, :3] clusters = [] st_time = time.time() kmeans = MiniBatchKMeans(n_clusters=root_num[0][0], random_state=0).fit(p) # print('Time on kmeans', time.time() - st_time) clusters.append([[kmeans.labels_]]) # centers = kmeans.cluster_centers_ for j in range(time_step): shape_states = calc_shape_states_RiceGrip(j * dt, dt, shape_state_dim, gripper_config) pyflex.set_shape_states(shape_states) positions[j, :n_particles, :3] = pyflex.get_rigidGlobalPositions().reshape(-1, 3) positions[j, :n_particles, 3:] = pyflex.get_positions().reshape(-1, 4)[:, :3] shape_states = pyflex.get_shape_states().reshape(-1, shape_state_dim) for k in range(n_shapes): positions[j, n_particles + k, :3] = shape_states[k, :3] positions[j, n_particles + k, 3:] = shape_states[k, :3] shape_quats[j, k] = shape_states[k, 6:10] if j > 0: velocities[j] = (positions[j] - positions[j - 1]) / dt pyflex.step() data = [positions[j], velocities[j], shape_quats[j], clusters, scene_params] store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5')) else: raise AssertionError("Unsupported env") # change dtype for more accurate stat calculation # only normalize positions and velocities datas = [positions.astype(np.float64), velocities.astype(np.float64)] # NOTE: stats is of length 2, for positions and velocities for j in range(len(stats)): stat = init_stat(stats[j].shape[0]) stat[:, 0] = np.mean(datas[j], axis=(0, 1))[:] stat[:, 1] = np.std(datas[j], axis=(0, 1))[:] stat[:, 2] = datas[j].shape[0] * datas[j].shape[1] stats[j] = combine_stat(stats[j], stat) pyflex.clean() return stats
def gen_PyFleX(info): env, env_idx = info['env'], info['env_idx'] thread_idx, data_dir, data_names = info['thread_idx'], info['data_dir'], info['data_names'] n_rollout, time_step = info['n_rollout'], info['time_step'] shape_state_dim, dt = info['shape_state_dim'], info['dt'] gen_vision = info['gen_vision'] vision_dir, vis_width, vis_height = info['vision_dir'], info['vis_width'], info['vis_height'] np.random.seed(round(time.time() * 1000 + thread_idx) % 2 ** 32) # positions stats = [init_stat(3)] import pyflex pyflex.init() for i in range(n_rollout): if i % 10 == 0: print("%d / %d" % (i, n_rollout)) rollout_idx = thread_idx * n_rollout + i rollout_dir = os.path.join(data_dir, str(rollout_idx)) os.system('mkdir -p ' + rollout_dir) if env == 'RigidFall': g_low, g_high = info['physics_param_range'] gravity = rand_float(g_low, g_high) print("Generated RigidFall rollout {} with gravity {} from range {} ~ {}".format( i, gravity, g_low, g_high)) n_instance = 3 draw_mesh = 1 scene_params = np.zeros(n_instance * 3 + 3) scene_params[0] = n_instance scene_params[1] = gravity scene_params[-1] = draw_mesh low_bound = 0.09 for j in range(n_instance): x = rand_float(0., 0.1) y = rand_float(low_bound, low_bound + 0.01) z = rand_float(0., 0.1) scene_params[j * 3 + 2] = x scene_params[j * 3 + 3] = y scene_params[j * 3 + 4] = z low_bound += 0.21 pyflex.set_scene(env_idx, scene_params, thread_idx) pyflex.set_camPos(np.array([0.2, 0.875, 2.0])) n_particles = pyflex.get_n_particles() n_shapes = 1 # the floor positions = np.zeros((time_step, n_particles + n_shapes, 3), dtype=np.float32) shape_quats = np.zeros((time_step, n_shapes, 4), dtype=np.float32) for j in range(time_step): positions[j, :n_particles] = pyflex.get_positions().reshape(-1, 4)[:, :3] ref_positions = positions[0] for k in range(n_instance): XX = ref_positions[64*k:64*(k+1)] YY = positions[j, 64*k:64*(k+1)] X = XX.copy().T Y = YY.copy().T mean_X = np.mean(X, 1, keepdims=True) mean_Y = np.mean(Y, 1, keepdims=True) X = X - mean_X Y = Y - mean_Y C = np.dot(X, Y.T) U, S, Vt = np.linalg.svd(C) D = np.eye(3) D[2, 2] = np.linalg.det(np.dot(Vt.T, U.T)) R = np.dot(Vt.T, np.dot(D, U.T)) t = mean_Y - np.dot(R, mean_X) YY_fitted = (np.dot(R, XX.T) + t).T # print("MSE fit", np.mean(np.square(YY_fitted - YY))) positions[j, 64*k:64*(k+1)] = YY_fitted if gen_vision: pyflex.step(capture=True, path=os.path.join(rollout_dir, str(j) + '.tga')) else: pyflex.step() data = [positions[j], shape_quats[j], scene_params] store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5')) if gen_vision: images = np.zeros((time_step, vis_height, vis_width, 3), dtype=np.uint8) for j in range(time_step): img_path = os.path.join(rollout_dir, str(j) + '.tga') img = scipy.misc.imread(img_path)[:, :, :3][:, :, ::-1] img = cv2.resize(img, (vis_width, vis_height), interpolation=cv2.INTER_AREA) images[j] = img os.system('rm ' + img_path) store_data(['positions', 'images', 'scene_params'], [positions, images, scene_params], os.path.join(vision_dir, str(rollout_idx) + '.h5')) elif env == 'MassRope': s_low, s_high = info['physics_param_range'] stiffness = rand_float(s_low, s_high) print("Generated MassRope rollout {} with gravity {} from range {} ~ {}".format( i, stiffness, s_low, s_high)) x = 0. y = 1.0 z = 0. length = 0.7 draw_mesh = 1. scene_params = np.array([x, y, z, length, stiffness, draw_mesh]) pyflex.set_scene(env_idx, scene_params, 0) pyflex.set_camPos(np.array([0.13, 2.0, 3.2])) action = np.zeros(3) # the last particle is the pin, regarded as shape n_particles = pyflex.get_n_particles() - 1 n_shapes = 1 # the mass at the top of the rope positions = np.zeros((time_step + 1, n_particles + n_shapes, 3), dtype=np.float32) shape_quats = np.zeros((time_step + 1, n_shapes, 4), dtype=np.float32) action = np.zeros(3) for j in range(time_step + 1): positions[j] = pyflex.get_positions().reshape(-1, 4)[:, :3] if j >= 1: # append the action (position of the pin) to the previous time step positions[j - 1, -1, :] = positions[j, -1, :] ref_positions = positions[0] # apply rigid projection to the rigid object # cube: [0, 81) # rope: [81, 95) # pin: [95, 96) XX = ref_positions[:81] YY = positions[j, :81] X = XX.copy().T Y = YY.copy().T mean_X = np.mean(X, 1, keepdims=True) mean_Y = np.mean(Y, 1, keepdims=True) X = X - mean_X Y = Y - mean_Y C = np.dot(X, Y.T) U, S, Vt = np.linalg.svd(C) D = np.eye(3) D[2, 2] = np.linalg.det(np.dot(Vt.T, U.T)) R = np.dot(Vt.T, np.dot(D, U.T)) t = mean_Y - np.dot(R, mean_X) YY_fitted = (np.dot(R, XX.T) + t).T positions[j, :81] = YY_fitted scale = 0.1 action[0] += rand_float(-scale, scale) - positions[j, -1, 0] * 0.1 action[2] += rand_float(-scale, scale) - positions[j, -1, 2] * 0.1 if gen_vision: pyflex.step(action * dt, capture=True, path=os.path.join(rollout_dir, str(j) + '.tga')) else: pyflex.step(action * dt) if j >= 1: data = [positions[j - 1], shape_quats[j - 1], scene_params] store_data(data_names, data, os.path.join(rollout_dir, str(j - 1) + '.h5')) if gen_vision: images = np.zeros((time_step, vis_height, vis_width, 3), dtype=np.uint8) for j in range(time_step): img_path = os.path.join(rollout_dir, str(j) + '.tga') img = scipy.misc.imread(img_path)[:, :, :3][:, :, ::-1] img = cv2.resize(img, (vis_width, vis_height), interpolation=cv2.INTER_AREA) images[j] = img os.system('rm ' + img_path) store_data(['positions', 'images', 'scene_params'], [positions, images, scene_params], os.path.join(vision_dir, str(rollout_idx) + '.h5')) else: raise AssertionError("Unsupported env") # change dtype for more accurate stat calculation # only normalize positions datas = [positions[:time_step].astype(np.float64)] for j in range(len(stats)): stat = init_stat(stats[j].shape[0]) stat[:, 0] = np.mean(datas[j], axis=(0, 1))[:] stat[:, 1] = np.std(datas[j], axis=(0, 1))[:] stat[:, 2] = datas[j].shape[0] * datas[j].shape[1] stats[j] = combine_stat(stats[j], stat) pyflex.clean() return stats
def reset(self, state): for i in range(100): pyflex.step()