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)
