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 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 _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_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 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 _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 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, 'glass_x': self.glass_x, 'glass_y': self.glass_y, 'glass_rotation': self.glass_rotation, 'glass_states': self.glass_states, 'poured_glass_states': self.poured_glass_states, 'glass_params': self.glass_params, 'config_id': self.current_config_id, 'line_box_x': self.line_box_x, 'line_box_y': self.line_box_y }
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 set_picker_pos(picker_pos): """ Caution! Should only be called during the reset of the environment. Used only for cloth drop environment. """ shape_states = np.array(pyflex.get_shape_states()).reshape(-1, 14) shape_states[0:2, 3:6] = picker_pos # picker 1 shape_states[0:2, :3] = picker_pos # picker 2 print("Posicion del picker" + str(picker_pos)) # Define hanger halfEdge = np.array([0.0225, 0.01, 0.165]) # height, width, depth center = np.array([0., 0., 0.]) quat = np.array([0., 0., -0., 1.]) pyflex.add_box(halfEdge, center, quat) # Add hanger position - REMOVE MAGIC NUMBERS wall_pos = [ -0.1275, 0.5, 0.0, -0.1275, 0.5, 0.0, 0., 0., -0., 1., 0., 0., -0., 1. ] shape_states = np.vstack([shape_states, wall_pos]) pyflex.set_shape_states(shape_states)
x_box = x_center v_box = 0 # simulation for i 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(x_box, x_box_last, scene_params[-2:]) pyflex.set_shape_states(shape_states_) positions[i] = pyflex.get_positions().reshape(-1, dim_position) velocities[i] = pyflex.get_velocities().reshape(-1, dim_velocity) shape_states[i] = pyflex.get_shape_states().reshape(-1, dim_shape_state) if i == 0: print(np.min(positions[i], 0), np.max(positions[i], 0)) print(x_box, box_dis_x, box_dis_z) pyflex.step() # playback pyflex.set_scene(6, scene_params, 0) for i in range(len(boxes) - 1): halfEdge = boxes[i][0] center = boxes[i][1] quat = boxes[i][2] pyflex.add_box(halfEdge, center, quat)
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 _set_pos(picker_pos, particle_pos): shape_states = np.array(pyflex.get_shape_states()).reshape(-1, 14) shape_states[:, 3:6] = shape_states[:, :3] shape_states[:, :3] = picker_pos pyflex.set_shape_states(shape_states) pyflex.set_positions(particle_pos)
def _get_pos(): """ Get the current pos of the pickers and the particles, along with the inverse mass of each particle """ picker_pos = np.array(pyflex.get_shape_states()).reshape(-1, 14) particle_pos = np.array(pyflex.get_positions()).reshape(-1, 4) return picker_pos[:, :3], particle_pos
def set_picker_pos(picker_pos): """ Caution! Should only be called during the reset of the environment. Used only for cloth drop environment. """ shape_states = np.array(pyflex.get_shape_states()).reshape(-1, 14) shape_states[:, 3:6] = picker_pos shape_states[:, :3] = picker_pos pyflex.set_shape_states(shape_states)
def get_picker_pos(self): picker_pos = np.array(pyflex.get_shape_states()).reshape(-1, 14) return picker_pos[:, :3]
# particles = np.load(path) # load_pos(particles) # pyflex.set_positions(pos) for i in range(time_step): # Load in the particles path = particle_path_prefix + "frame_{}.npy".format(i) particles = np.load(path) / scale clear_pos() load_pos(particles) pyflex.set_positions(pos) # pos[0] = np.array([i/100, 0.3, i/100, 1]) # pyflex.set_positions(pos) # pos = pyflex.get_positions() # print(pos) # pyflex.step(capture=1, path=os.path.join(des_dir, "render_%04d.tga" % i)) # if i == 200: path = tool_path_prefix + "frame_{}.npy".format(i) tool = np.load(path) # print(tool) # break states = pyflex.get_shape_states() # # print(states) states[17] = (tool[0] + 0.5) / scale states[18] = tool[1] / scale + 0.15 states[19] = tool[2] / scale pyflex.set_shape_states(states) pyflex.render(capture=1, path=os.path.join(des_dir, "render_%04d.tga" % i)) pyflex.clean()