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
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) for i in range(time_step): pyflex.set_positions(positions[i]) pyflex.set_shape_states(shape_states[i, :-1]) pyflex.render(capture=1, path=os.path.join(des_dir, 'render_%d.tga' % i)) pyflex.clean()
def close(self): pyflex.clean()
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_Cloth(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'] dt, args, phase = info['dt'], info['args'], info['phase'] vis_width, vis_height = info['vis_width'], info['vis_height'] state_dim = args.state_dim action_dim = args.action_dim dt = 1. / 60. np.random.seed(round(time.time() * 1000 + thread_idx) % 2 ** 32) stats = [init_stat(state_dim), init_stat(action_dim)] engine = ClothEngine(dt, state_dim, action_dim) import pyflex pyflex.init() # bar = ProgressBar() for i in range(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) engine.init(pyflex) scene_params = engine.scene_params action = np.zeros(4) states_all = np.zeros((time_step, engine.n_particles, state_dim)) actions_all = np.zeros((time_step, 1, action_dim)) # drop the cloth down engine.set_action(action) engine.step() for j in range(time_step): positions = pyflex.get_positions().reshape(-1, 4)[:, :3] # sample the action if j % 5 == 0: ctrl_pts = rand_int(0, 8) act_lim = 0.05 dx = rand_float(-act_lim, act_lim) dz = rand_float(-act_lim, act_lim) dy = 0.05 action = np.array([ctrl_pts, dx, dy, dz]) else: action[2] = 0. # store the rollout information state = engine.get_state() states_all[j] = state tga_path = os.path.join(rollout_dir, '%d.tga' % j) pyflex.render(capture=True, path=tga_path) tga = Image.open(tga_path) tga = np.array(tga)[:, 60:780, :3][:, :, ::-1] tga = cv2.resize(tga, (vis_width, vis_height), interpolation=cv2.INTER_AREA) os.system('rm ' + tga_path) jpg_path = os.path.join(rollout_dir, 'fig_%d.jpg' % j) cv2.imwrite(jpg_path, tga) actions_all[j, 0] = action.copy() engine.set_action(action) engine.step() datas = [states_all, actions_all, scene_params] store_data(data_names, datas, rollout_dir + '.h5') datas = [datas[j].astype(np.float64) for j in range(len(datas))] 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] stats[j] = combine_stat(stats[j], stat) pyflex.clean() return stats