def set_state(self, state_dict): pyflex.set_positions(state_dict['particle_pos']) pyflex.set_velocities(state_dict['particle_vel']) pyflex.set_shape_states(state_dict['shape_pos']) pyflex.set_phases(state_dict['phase']) self.camera_params = copy.deepcopy(state_dict['camera_params']) self.update_camera(self.camera_name)
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 set_shape_states(self, glass_states, poured_glass_states): ''' set the the shape states of both glasses. ''' all_states = np.concatenate((glass_states, poured_glass_states), axis=0) pyflex.set_shape_states(all_states)
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 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 _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 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)
# 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()
R[1, 0] = 2 * b * c + 2 * a * d R[1, 1] = a**2 - b**2 + c**2 - d**2 R[1, 2] = 2 * c * d - 2 * a * b R[2, 0] = 2 * b * d - 2 * a * c R[2, 1] = 2 * c * d + 2 * a * b R[2, 2] = a**2 - b**2 - c**2 + d**2 return np.dot(R, p) for r in range(grip_time): gripper_config = sample_gripper_config() # gripper_config = saved[r] # print(gripper_config) for i in range(time_step): shape_states_ = calc_shape_states(i * dt, gripper_config) pyflex.set_shape_states(shape_states_) positions[r, i] = pyflex.get_positions().reshape(-1, dim_position) if i == 0: surface_idx = calc_surface_idx(positions[r, i, :, :3]) # visualize_point_cloud(positions[i, :, :3], surface_idx) restPositions[r, i] = pyflex.get_restPositions().reshape(-1, dim_position) velocities[r, i] = pyflex.get_velocities().reshape(-1, dim_velocity) shape_states[r, i] = pyflex.get_shape_states().reshape(-1, dim_shape_state) rigid_offsets[r, i] = pyflex.get_rigidOffsets().reshape(-1, 1) rigid_indices[r, i] = pyflex.get_rigidIndices().reshape(-1, 1) rigid_localPositions[r, i] = pyflex.get_rigidLocalPositions().reshape(-1, 3) rigid_globalPositions[r, i] = pyflex.get_rigidGlobalPositions().reshape(-1, 3)
positions = np.zeros((time_step, n_particles, dim_position)) velocities = np.zeros((time_step, n_particles, dim_velocity)) shape_states = np.zeros((time_step, n_shapes, dim_shape_state)) 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):
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_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 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 glass params if states is None: self.set_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_params = glass_params # create glass self.create_glass(self.glass_dis_x, self.glass_dis_z, self.height, self.border) # move glass to be at ground or on the table self.glass_states = self.init_glass_state(self.x_center, 0, self.glass_dis_x, self.glass_dis_z, self.height, self.border) pyflex.set_shape_states(self.glass_states) # record glass floor center x self.glass_x = self.x_center # 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.2 * fluid_radius, fluid_radius * 0.45, 1.2 * fluid_radius]) 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() # pyflex.render() # time.sleep(0.1) 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, return_sum=False) 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(100): pyflex.step() # pyflex.render() # time.sleep(0.1) 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, return_sum=False) not_in_glass = 1 - in_glass not_total_num = np.sum(not_in_glass) for _ in range(30): pyflex.step() # pyflex.render() else: # set to passed-in cached init states self.set_state(states)
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)
elif args.env == 'FluidShake': x, y, z, dim_x, dim_y, dim_z, box_dis_x, box_dis_z = scene_params boxes = calc_box_init_FluidShake(box_dis_x, box_dis_z, height, border) x_box = x + (dim_x-1)/2.*0.055 for box_idx in range(len(boxes) - 1): halfEdge = boxes[box_idx][0] center = boxes[box_idx][1] quat = boxes[box_idx][2] pyflex.add_box(halfEdge, center, quat) for step in range(args.time_step - 1): if args.env == 'RiceGrip': pyflex.set_shape_states(s_gt[step]) elif args.env == 'FluidShake': pyflex.set_shape_states(s_gt[step, :-1]) mass = np.zeros((n_particles, 1)) if args.env == 'RiceGrip': p = np.concatenate([p_gt[step, :n_particles, -3:], mass], 1) else: p = np.concatenate([p_gt[step, :n_particles], mass], 1) pyflex.set_positions(p) pyflex.render(capture=1, path=os.path.join(des_dir, 'gt_%d.tga' % step)) ##### render for the predictions pyflex.set_scene(env_idx, scene_params, 0)