def render(self, mode='rgb_array'):
if mode == 'rgb_array':
img, depth = pyflex.render()
width, height = self.camera_params['default_camera']['width'], self.camera_params['default_camera']['height']
img = img.reshape(height, width, 4)[::-1, :, :3] # Need to reverse the height dimension
return img
elif mode == 'human':
raise NotImplementedError
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 show_depth():
# render rgb and depth
img, depth = pyflex.render()
img = img.reshape((720, 720, 4))[::-1, :, :3]
depth = depth.reshape((720, 720))[::-1]
# get foreground mask
rgb, depth = pyflex.render_cloth()
depth = depth.reshape(720, 720)[::-1]
# mask = mask[:, :, 3]
# depth[mask == 0] = 0
# show rgb and depth(masked)
depth[depth > 5] = 0
fig, axes = plt.subplots(1, 2, figsize=(12, 5))
axes[0].imshow(img)
axes[1].imshow(depth)
plt.show()
def _get_center_point(self, pos):
pos = np.reshape(pos, [-1, 4])
min_x = np.min(pos[:, 0])
min_y = np.min(pos[:, 2])
max_x = np.max(pos[:, 0])
max_y = np.max(pos[:, 2])
return 0.5 * (min_x + max_x), 0.5 * (min_y + max_y)
if __name__ == '__main__':
env = RopeNewEnv(observation_mode='key_point',
action_mode='picker',
num_picker=2,
render=True,
headless=False,
horizon=75,
action_repeat=8,
num_variations=10,
use_cached_states=False,
save_cached_states=False,
deterministic=False)
env.reset(config=env.get_default_config())
for i in range(1000):
print(i)
print("right before pyflex step")
pyflex.step()
print("right after pyflex step")
print("right before pyflex render")
pyflex.render()
print("right after pyflex render")
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 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
print(pyflex.get_sceneParams())
pyflex.step()
for i in range(time_step - 1):
idx = 10
cnt = 0
for j in range(n_rigids):
st, ed = rigid_offsets[r, i, j, 0], rigid_offsets[r, i, j + 1, 0]
for k in range(st, ed):
if rigid_indices[r, i, cnt, 0] == idx:
print(i, j, positions[r, i, rigid_indices[r, i, cnt, 0], :3],
rigid_globalPositions[r, i, rigid_indices[r, i, cnt, 0]],
rotate(rigid_localPositions[r, i, cnt], rigid_rotations[r, i, j]) + \
rigid_translations[r, i, j])
cnt += 1
pyflex.set_scene(5, scene_params, 0)
pyflex.add_box(halfEdge, center, quat)
pyflex.add_box(halfEdge, center, quat)
for r in range(grip_time):
for i in range(time_step):
pyflex.set_positions(positions[r, i])
pyflex.set_shape_states(shape_states[r, i])
pyflex.render(capture=1, path=os.path.join(des_dir, 'render_%d.tga' % (r * time_step + i)))
pyflex.clean()
# Set all the particles to unseen locations
def clear_pos():
pos[:, 0:3] = np.ones((num_particles, 3)) * 20
# Function for putting particle positions into the array
def load_pos(particles):
print(particles.shape[0])
for i in range(particles.shape[0]):
if i < num_particles:
pos[i, 0:3] = particles[i, 0:3]
# # Load in the particles
# path = particle_path_prefix+"frame_{}.npy".format(599)
# 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)
# pyflex.step(capture=1, path=os.path.join(des_dir, "render_%04d.tga" % i))
pyflex.render(capture=1, path=os.path.join(des_dir, "render_%04d.tga" % i))
pyflex.clean()
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)
if args.env == 'RiceGrip':
pyflex.add_box(halfEdge, center, quat)
pyflex.add_box(halfEdge, center, quat)
elif args.env == 'FluidShake':
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):
