boxes = calc_box_init(box_dis_x, box_dis_z)
for i in range(len(boxes)):
halfEdge = boxes[i][0]
center = boxes[i][1]
quat = boxes[i][2]
pyflex.add_box(halfEdge, center, quat)
### read scene info
print("Scene Upper:", pyflex.get_scene_upper())
print("Scene Lower:", pyflex.get_scene_lower())
print("Num particles:", pyflex.get_phases().reshape(-1, 1).shape[0])
print("Phases:", np.unique(pyflex.get_phases()))
n_particles = pyflex.get_n_particles()
n_shapes = pyflex.get_n_shapes()
n_rigids = pyflex.get_n_rigids()
n_rigidPositions = pyflex.get_n_rigidPositions()
print("n_particles", n_particles)
print("n_shapes", n_shapes)
print("n_rigids", n_rigids)
print("n_rigidPositions", n_rigidPositions)
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
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