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 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 generate_env_variation(self, num_variations=2, vary_cloth_size=True):
""" Generate initial states. Note: This will also change the current states! """
max_wait_step = 1000 # Maximum number of steps waiting for the cloth to stablize
stable_vel_threshold = 0.2 # Cloth stable when all particles' vel are smaller than this
generated_configs, generated_states = [], []
default_config = self.get_default_config()
default_config['flip_mesh'] = 1
for i in range(num_variations):
config = deepcopy(default_config)
self.update_camera(config['camera_name'],
config['camera_params'][config['camera_name']])
if vary_cloth_size:
cloth_dimx, cloth_dimy = self._sample_cloth_size()
config['ClothSize'] = [cloth_dimx, cloth_dimy]
else:
cloth_dimx, cloth_dimy = config['ClothSize']
self.set_scene(config)
self.action_tool.reset([0., -1., 0.])
pos = pyflex.get_positions().reshape(-1, 4)
pos[:, :3] -= np.mean(pos, axis=0)[:3]
if self.action_mode in ['sawyer', 'franka'
]: # Take care of the table in robot case
pos[:, 1] = 0.57
else:
pos[:, 1] = 0.005
pos[:, 3] = 1
pyflex.set_positions(pos.flatten())
pyflex.set_velocities(np.zeros_like(pos))
for _ in range(5): # In case if the cloth starts in the air
pyflex.step()
for wait_i in range(max_wait_step):
pyflex.step()
curr_vel = pyflex.get_velocities()
if np.alltrue(np.abs(curr_vel) < stable_vel_threshold):
break
center_object()
angle = (np.random.random() - 0.5) * np.pi / 2
self.rotate_particles(angle)
generated_configs.append(deepcopy(config))
print('config {}: {}'.format(i, config['camera_params']))
generated_states.append(deepcopy(self.get_state()))
return generated_configs, generated_states
def generate_env_variation(self, num_variations=1, vary_cloth_size=True):
""" Generate initial states. Note: This will also change the current states! """
max_wait_step = 300 # Maximum number of steps waiting for the cloth to stablize
stable_vel_threshold = 0.01 # Cloth stable when all particles' vel are smaller than this
generated_configs, generated_states = [], []
default_config = self.get_default_config()
for i in range(num_variations):
config = deepcopy(default_config)
self.update_camera(config['camera_name'],
config['camera_params'][config['camera_name']])
if vary_cloth_size:
cloth_dimx, cloth_dimy = self._sample_cloth_size()
config['ClothSize'] = [cloth_dimx, cloth_dimy]
else:
cloth_dimx, cloth_dimy = config['ClothSize']
self.set_scene(config)
self.action_tool.reset([0., -1., 0.])
pos = pyflex.get_positions().reshape(-1, 4)
pos[:, :3] -= np.mean(pos, axis=0)[:3]
if self.action_mode in ['sawyer', 'franka'
]: # Take care of the table in robot case
pos[:, 1] = 0.57
else:
pos[:, 1] = 0.005
pos[:, 3] = 1
pyflex.set_positions(pos.flatten())
pyflex.set_velocities(np.zeros_like(pos))
pyflex.step()
num_particle = cloth_dimx * cloth_dimy
pickpoint = random.randint(0, num_particle - 1)
curr_pos = pyflex.get_positions()
original_inv_mass = curr_pos[pickpoint * 4 + 3]
curr_pos[
pickpoint * 4 +
3] = 0 # Set the mass of the pickup point to infinity so that it generates enough force to the rest of the cloth
pickpoint_pos = curr_pos[pickpoint * 4:pickpoint * 4 + 3].copy(
) # Pos of the pickup point is fixed to this point
pickpoint_pos[1] += np.random.random(1) * 0.5 + 0.5
pyflex.set_positions(curr_pos)
# Pick up the cloth and wait to stablize
for j in range(0, max_wait_step):
curr_pos = pyflex.get_positions()
curr_vel = pyflex.get_velocities()
curr_pos[pickpoint * 4:pickpoint * 4 + 3] = pickpoint_pos
curr_vel[pickpoint * 3:pickpoint * 3 + 3] = [0, 0, 0]
pyflex.set_positions(curr_pos)
pyflex.set_velocities(curr_vel)
pyflex.step()
if np.alltrue(
np.abs(curr_vel) < stable_vel_threshold) and j > 5:
break
# Drop the cloth and wait to stablize
curr_pos = pyflex.get_positions()
curr_pos[pickpoint * 4 + 3] = original_inv_mass
pyflex.set_positions(curr_pos)
for _ in range(max_wait_step):
pyflex.step()
curr_vel = pyflex.get_velocities()
if np.alltrue(curr_vel < stable_vel_threshold):
break
center_object()
if self.action_mode == 'sphere' or self.action_mode.startswith(
'picker'):
curr_pos = pyflex.get_positions()
self.action_tool.reset(curr_pos[pickpoint * 4:pickpoint * 4 +
3] + [0., 0.2, 0.])
generated_configs.append(deepcopy(config))
generated_states.append(deepcopy(self.get_state()))
self.current_config = config # Needed in _set_to_flatten function
generated_configs[-1]['flatten_area'] = self._set_to_flatten(
) # Record the maximum flatten area
print('config {}: camera params {}, flatten area: {}'.format(
i, config['camera_params'],
generated_configs[-1]['flatten_area']))
return generated_configs, generated_states