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 = 500 # Maximum number of steps waiting for the cloth to stablize
stable_vel_threshold = 0.1 # 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.])
pickpoints = self._get_drop_point_idx(
)[:2] # Pick two corners of the cloth and wait until stablize
config['target_pos'] = self._get_flat_pos()
self._set_to_vertical(x_low=np.random.random() * 0.2 - 0.1,
height_low=np.random.random() * 0.1 + 0.1)
# Get height of the cloth without the gravity. With gravity, it will be longer
p1, _, p2, _ = self._get_key_point_idx()
curr_pos = pyflex.get_positions().reshape(-1, 4)
curr_pos[0] += np.random.random() * 0.001 # Add small jittering
original_inv_mass = curr_pos[pickpoints, 3]
curr_pos[
pickpoints,
3] = 0 # Set mass of the pickup point to infinity so that it generates enough force to the rest of the cloth
pickpoint_pos = curr_pos[pickpoints, :3]
pyflex.set_positions(curr_pos.flatten())
picker_radius = self.action_tool.picker_radius
self.action_tool.update_picker_boundary([-0.3, 0.05, -0.5],
[0.5, 2, 0.5])
self.action_tool.set_picker_pos(picker_pos=pickpoint_pos +
np.array([0., picker_radius, 0.]))
# Pick up the cloth and wait to stablize
for j in range(0, max_wait_step):
pyflex.step()
curr_pos = pyflex.get_positions().reshape((-1, 4))
curr_vel = pyflex.get_velocities().reshape((-1, 3))
if np.alltrue(curr_vel < stable_vel_threshold) and j > 300:
break
curr_pos[pickpoints, :3] = pickpoint_pos
pyflex.set_positions(curr_pos)
curr_pos = pyflex.get_positions().reshape((-1, 4))
curr_pos[pickpoints, 3] = original_inv_mass
pyflex.set_positions(curr_pos.flatten())
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 random_pick_and_place(pick_num=10, pick_scale=0.01):
""" Random pick a particle up and the drop it for pick_num times"""
curr_pos = pyflex.get_positions().reshape(-1, 4)
num_particles = curr_pos.shape[0]
for i in range(pick_num):
pick_id = np.random.randint(num_particles)
pick_dir = np.random.random(3) * 2 - 1
pick_dir[1] = (pick_dir[1] + 1)
pick_dir *= pick_scale
original_inv_mass = curr_pos[pick_id, 3]
for _ in range(60):
curr_pos = pyflex.get_positions().reshape(-1, 4)
curr_pos[pick_id, :3] += pick_dir
curr_pos[pick_id, 3] = 0
pyflex.set_positions(curr_pos.flatten())
pyflex.step()
# Revert mass
curr_pos = pyflex.get_positions().reshape(-1, 4)
curr_pos[pick_id, 3] = original_inv_mass
pyflex.set_positions(curr_pos.flatten())
pyflex.step()
# Wait to stabalize
for _ in range(100):
pyflex.step()
curr_vel = pyflex.get_velocities()
if np.alltrue(curr_vel < 0.01):
break
for _ in range(500):
pyflex.step()
curr_vel = pyflex.get_velocities()
if np.alltrue(curr_vel < 0.01):
break
def _get_obs(self):
'''
return the observation based on the current flex state.
'''
if self.observation_mode == 'cam_rgb':
return self.get_image(self.camera_width, self.camera_height)
elif self.observation_mode in ['point_cloud', 'key_point']:
if self.observation_mode == 'point_cloud':
particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3].flatten()
pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float)
pos[:len(particle_pos)] = particle_pos
else:
pos = np.empty(0, dtype=np.float)
particle_state = pyflex.get_positions().reshape((-1, self.dim_position))
torus_center_y = np.mean(particle_state[:, 1])
torus_center_x = np.mean(particle_state[:, 0])
on_box = float(torus_center_y >= self.height)
within_box_bound = float(torus_center_x > self.box_x - 0.5 * self.box_dis_x and
torus_center_x < self.box_x + self.box_dis_x * 0.5)
cup_state = np.array([self.box_x, self.box_dis_x, self.box_dis_z, self.height,
torus_center_x, torus_center_y, on_box, within_box_bound])
return np.hstack([pos, cup_state]).flatten()
else:
raise NotImplementedError
def _get_obs(self):
if self.observation_mode == 'cam_rgb':
obs_img = self.get_image(self.camera_height, self.camera_width)
goal_img = self.current_config['goal_character_img']
ret_img = np.concatenate([obs_img, goal_img], axis=2)
return ret_img
if self.observation_mode == 'point_cloud':
particle_pos = np.array(pyflex.get_positions()).reshape(
[-1, 4])[:, :3].flatten()
pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float)
pos[:len(particle_pos)] = particle_pos
pos[len(particle_pos):] = self.current_config[
"goal_character_pos"][:, :3].flatten()
elif self.observation_mode == 'key_point':
particle_pos = np.array(pyflex.get_positions()).reshape([-1,
4])[:, :3]
keypoint_pos = particle_pos[self.key_point_indices, :3]
goal_keypoint_pos = self.current_config["goal_character_pos"][
self.key_point_indices, :3]
pos = np.concatenate([keypoint_pos, goal_keypoint_pos],
axis=0).flatten()
if self.action_mode in ['sphere', 'picker']:
shapes = pyflex.get_shape_states()
shapes = np.reshape(shapes, [-1, 14])
pos = np.concatenate([pos.flatten(), shapes[:, :3].flatten()])
return pos
def _get_obs(self):
'''
return the observation based on the current flex state.
'''
if self.observation_mode == 'cam_rgb':
return self.get_image(self.camera_width, self.camera_height)
elif self.observation_mode == 'point_cloud':
particle_pos = np.array(pyflex.get_positions()).reshape(
[-1, 4])[:, :3].flatten()
pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float)
pos[:len(particle_pos)] = particle_pos
return pos.flatten()
elif self.observation_mode == 'key_point':
pos = np.empty(0, dtype=np.float)
water_state = pyflex.get_positions().reshape([-1, 4])
water_num = len(water_state)
in_glass = self.in_glass(water_state, self.glass_states,
self.border, self.height)
out_glass = water_num - in_glass
in_glass = float(in_glass) / water_num
out_glass = float(out_glass) / water_num
cup_state = np.array([
self.glass_x, self.glass_dis_x, self.glass_dis_z, self.height,
self._get_current_water_height(), in_glass, out_glass
])
return np.hstack([pos, cup_state]).flatten()
else:
raise NotImplementedError
def _reset(self):
""" Right now only use one initial state. Need to make sure _reset always give the same result. Otherwise CEM will fail."""
if hasattr(self, 'action_tool'):
particle_pos = pyflex.get_positions().reshape(-1, 4)
p1, p2, p3, p4 = self._get_key_point_idx()
key_point_pos = particle_pos[(
p1, p2), :3] # Was changed from from p1, p4.
middle_point = np.mean(key_point_pos, axis=0)
self.action_tool.reset([middle_point[0], 0.1, middle_point[2]])
# picker_low = self.action_tool.picker_low
# picker_high = self.action_tool.picker_high
# offset_x = self.action_tool._get_pos()[0][0][0] - picker_low[0] - 0.3
# picker_low[0] += offset_x
# picker_high[0] += offset_x
# picker_high[0] += 1.0
# self.action_tool.update_picker_boundary(picker_low, picker_high)
config = self.get_current_config()
num_particles = np.prod(config['ClothSize'], dtype=int)
particle_grid_idx = np.array(list(range(num_particles))).reshape(
config['ClothSize'][1],
config['ClothSize'][0]) # Reversed index here
cloth_dimx = config['ClothSize'][0]
x_split = cloth_dimx // 2
self.fold_group_a = particle_grid_idx[:, :x_split].flatten()
self.fold_group_b = np.flip(particle_grid_idx,
axis=1)[:, :x_split].flatten()
colors = np.zeros(num_particles)
colors[self.fold_group_a] = 1
# self.set_colors(colors) # TODO the phase actually changes the cloth dynamics so we do not change them for now. Maybe delete this later.
pyflex.step()
self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3]
pos_a = self.init_pos[self.fold_group_a, :]
pos_b = self.init_pos[self.fold_group_b, :]
self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1))
# If cached_goal is used recover goals from flatten initial states and set one for current iteration
if self.cached_goal:
# Get a random state from cached goals
rand_goal = np.random.randint(0, len(self.cached_goal_states))
self.goal_pos = self.cached_goal_states[rand_goal][
'particle_pos'].reshape((-1, 4))[:, :3]
#print("GOAL VS INIT", self.goal_pos.shape, self.init_pos.shape)
# Compute offset to match initial position TODO
# Trim goal_pos to match curr pos
if self.init_pos.shape[0] < self.goal_pos.shape[0]:
self.goal_pos = np.delete(
self.goal_pos,
np.s_[-(abs(self.goal_pos.shape[0] -
self.init_pos.shape[0]) + 1):-1], 0)
self.performance_init = None
info = self._get_info()
self.performance_init = info['performance']
return self._get_obs()
def _reset(self):
""" Right now only use one initial state"""
self.prev_covered_area = self._get_current_covered_area(
pyflex.get_positions())
if hasattr(self, 'action_tool'):
curr_pos = pyflex.get_positions()
cx, cy = self._get_center_point(curr_pos)
self.action_tool.reset([cx, 0.2, cy])
pyflex.step()
self.init_covered_area = None
info = self._get_info()
self.init_covered_area = info['performance']
return self._get_obs()
def _get_info(self):
goal_c_pos = self.current_config["goal_character_pos"][:, :3]
current_pos = pyflex.get_positions().reshape((-1, 4))[:, :3]
# way1: index matching
if self.reward_type == 'index':
dist = np.linalg.norm(current_pos - goal_c_pos, axis=1)
reward = -np.mean(dist)
if self.reward_type == 'bigraph':
# way2: downsample and then use Hungarian algorithm for bipartite graph matching
downsampled_cur_pos = current_pos[self.key_point_indices]
downsampled_goal_pos = goal_c_pos[self.key_point_indices]
W = np.zeros((len(downsampled_cur_pos), len(downsampled_cur_pos)))
for idx in range(len(downsampled_cur_pos)):
all_dist = np.linalg.norm(downsampled_cur_pos[idx] - downsampled_goal_pos, axis=1)
W[idx, :] = all_dist
row_idx, col_idx = opt.linear_sum_assignment(W)
dist = W[row_idx, col_idx].sum()
reward = -dist / len(downsampled_goal_pos)
performance = reward
performance_init = performance if self.performance_init is None else self.performance_init # Use the original performance
return {
'performance': performance,
'normalized_performance': (performance - performance_init) / (self.reward_max - performance_init),
}
def compute_reward(self, action=None, obs=None, set_prev_reward=False):
"""
The particles are splitted into two groups. The reward will be the minus average eculidean distance between each
particle in group a and the crresponding particle in group b
:param pos: nx4 matrix (x, y, z, inv_mass)
"""
# Get current pos
pos = pyflex.get_positions()
pos = pos.reshape((-1, 4))[:, :3]
if self.cached_goal:
# Trim pos to match goal cloth size if needed
if self.goal_pos.shape[0] < self.init_pos.shape[0]:
pos = np.delete(
pos, np.s_[-(abs(pos.shape[0] - self.goal_pos.shape[0]) +
1):-1], 0)
#curr_dist = (np.square(pos - self.goal_pos)).mean(axis=None)
curr_dist = mse(
self.goal_pos,
pos) # Calc min square error between goal and curr pos
else:
pos_group_a = pos[self.fold_group_a]
pos_group_b = pos[self.fold_group_b]
pos_group_b_init = self.init_pos[self.fold_group_b]
curr_dist = np.mean(np.linalg.norm(pos_group_a - pos_group_b, axis=1)) + \
1.2 * np.mean(np.linalg.norm(pos_group_b - pos_group_b_init, axis=1))
reward = -curr_dist
return reward
def _get_info(self):
# Duplicate of the compute reward function!
pos = pyflex.get_positions()
pos = pos.reshape((-1, 4))[:, :3]
pos_group_a = pos[self.fold_group_a]
pos_group_b = pos[self.fold_group_b]
pos_group_b_init = self.init_pos[self.fold_group_b]
group_dist = np.mean(np.linalg.norm(pos_group_a - pos_group_b, axis=1))
fixation_dist = np.mean(
np.linalg.norm(pos_group_b - pos_group_b_init, axis=1))
performance = -group_dist - 1.2 * fixation_dist
performance_init = performance if self.performance_init is None else self.performance_init # Use the original performance
info = {
'performance':
performance,
'normalized_performance':
(performance - performance_init) / (0. - performance_init),
'neg_group_dist':
-group_dist,
'neg_fixation_dist':
-fixation_dist,
'picker_pos':
self.get_picker_pos()
#'depth_map': self.get_depth_map()
}
if 'qpg' in self.action_mode:
info['total_steps'] = self.action_tool.total_steps
return info
def _reset(self):
""" Right now only use one initial state"""
if hasattr(self, 'action_tool'):
self.action_tool.reset([0., 0.2, 0.])
config = self.get_current_config()
self.flat_pos = self._get_flat_pos()
num_particles = np.prod(config['ClothSize'], dtype=int)
particle_grid_idx = np.array(list(range(num_particles))).reshape(config['ClothSize'][1], config['ClothSize'][0]) # Reversed index here
cloth_dimx = config['ClothSize'][0]
x_split = cloth_dimx // 2
self.fold_group_a = particle_grid_idx[:, :x_split].flatten()
self.fold_group_b = np.flip(particle_grid_idx, axis=1)[:, :x_split].flatten()
colors = np.zeros(num_particles)
colors[self.fold_group_a] = 1
pyflex.step()
self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3]
pos_a = self.init_pos[self.fold_group_a, :]
pos_b = self.init_pos[self.fold_group_b, :]
self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1))
self.performance_init = None
info = self._get_info()
self.performance_init = info['performance']
return self._get_obs()
def get_depth_map(self, pix_size=0.015625):
pos = pyflex.get_positions()
pos = pos.reshape((-1, 4))[:, :3]
dmap = np.zeros(shape=(self.camera_height, self.camera_width))
d_len_h = dmap.shape[0]
d_len_w = dmap.shape[1]
# Loop throught every cloth particle
for i in range(pos.shape[0]):
d_h = (pos[i][2] - 0.5)/pix_size # From position to pixel
d_w = (pos[i][0] - 0.5)/pix_size
d_h = int(d_len_h/2 + d_h)
d_w = int(d_len_w/2 + d_w)
# Assign correct depth to each pixel in dmap
if (abs(d_h) < d_len_h) and (abs(d_w) < d_len_w):
dmap[d_h][d_w] = pos[i][1]
# Recover also picker height
p_pos = self.get_picker_pos()[0]
u = (p_pos[2] - 0.5)/pix_size # From position to pixel
v = (p_pos[0] - 0.5)/pix_size
u = int(d_len_h/2 + u)
v = int(d_len_w/2 + v)
# Assign correct depth to each pixel in dmap
if (abs(u) < d_len_h) and (abs(v) < d_len_w):
dmap[u][v] = p_pos[1]
#img = Image.fromarray(dmap, 'I') # debug
#img.show()
return dmap
def _get_obs(self):
if self.observation_mode == 'cam_rgb':
return self.get_image(self.camera_height, self.camera_width)
if self.observation_mode == 'point_cloud':
particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3].flatten()
pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float)
pos[:len(particle_pos)] = particle_pos
elif self.observation_mode == 'key_point':
particle_pos = np.array(pyflex.get_positions()).reshape([-1, 4])[:, :3]
keypoint_pos = particle_pos[self._get_key_point_idx(), :3]
pos = keypoint_pos
if self.action_mode in ['sphere', 'picker', 'pickerTEO']:
shapes = pyflex.get_shape_states()
shapes = np.reshape(shapes, [-1, 14])
pos = np.concatenate([pos.flatten(), shapes[:, 0:3].flatten()])
return pos
def generate_alphabet_image(self):
self.goal_characters_image = {}
for c in self.goal_characters:
default_config = self.get_default_config(c=c)
goal_c_pos = self.goal_characters_position[c]
self.set_scene(default_config)
all_positions = pyflex.get_positions().reshape([-1, 4])
all_positions = goal_c_pos.copy() # ignore the first a few cloth particles
pyflex.set_positions(all_positions)
self.update_camera('default_camera', default_config['camera_params']['default_camera']) # why we need to do this?
self.action_tool.reset([0., -1., 0.]) # hide picker
# goal_c_img = self.get_image(self.camera_height, self.camera_width)
# import time
# for _ in range(50):
# pyflex.step(render=True)
# time.sleep(0.1)
# cv2.imshow('img', self.get_image())
# cv2.waitKey()
goal_c_img = self.get_image(self.camera_height, self.camera_width)
# cv2.imwrite('../data/images/rope-configuration-goal-image-{}.png'.format(c), goal_c_img[:,:,::-1])
# exit()
self.goal_characters_image[c] = goal_c_img.copy()
def compute_reward(self, action=None, obs=None, **kwargs):
""" Reward is the matching degree to the goal character"""
goal_c_pos = self.current_config["goal_character_pos"][:, :3]
current_pos = pyflex.get_positions().reshape((-1, 4))[:, :3]
# way1: index matching
if self.reward_type == 'index':
dist = np.linalg.norm(current_pos - goal_c_pos, axis=1)
reward = -np.mean(dist)
if self.reward_type == 'bigraph':
# way2: downsample and then use Hungarian algorithm for bipartite graph matching
downsampled_cur_pos = current_pos[self.key_point_indices]
downsampled_goal_pos = goal_c_pos[self.key_point_indices]
W = np.zeros((len(downsampled_cur_pos), len(downsampled_cur_pos)))
for idx in range(len(downsampled_cur_pos)):
all_dist = np.linalg.norm(downsampled_cur_pos[idx] -
downsampled_goal_pos,
axis=1)
W[idx, :] = all_dist
row_idx, col_idx = opt.linear_sum_assignment(W)
dist = W[row_idx, col_idx].sum()
reward = -dist / len(downsampled_goal_pos)
return reward
def _get_obs(self):
'''
return the observation based on the current flex state.
'''
if self.observation_mode == 'cam_rgb':
return self.get_image(self.camera_width, self.camera_height)
elif self.observation_mode in ['point_cloud', 'key_point']:
if self.observation_mode == 'point_cloud':
particle_pos = np.array(pyflex.get_positions()).reshape(
[-1, 4])[:, :3].flatten()
pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float)
pos[:len(particle_pos)] = particle_pos
cup_state = np.array([
self.glass_x, self.glass_y, self.glass_rotation,
self.glass_dis_x, self.glass_dis_z, self.height,
self.glass_distance + self.glass_x, self.poured_height,
self.poured_glass_dis_x, self.poured_glass_dis_z,
self.line_box_y, self.current_config['targe_amount']
])
else:
pos = np.empty(0, dtype=np.float)
cup_state = np.array([
self.glass_x, self.glass_y, self.glass_rotation,
self.glass_dis_x, self.glass_dis_z, self.height,
self.glass_distance + self.glass_x, self.poured_height,
self.poured_glass_dis_x, self.poured_glass_dis_z,
self._get_current_water_height(), self.line_box_y,
self.current_config['target_amount']
])
return np.hstack([pos, cup_state]).flatten()
else:
raise NotImplementedError
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 get_state(self):
pos = pyflex.get_positions()
vel = pyflex.get_velocities()
shape_pos = pyflex.get_shape_states()
phase = pyflex.get_phases()
camera_params = copy.deepcopy(self.camera_params)
return {'particle_pos': pos, 'particle_vel': vel, 'shape_pos': shape_pos, 'phase': phase, 'camera_params': camera_params,
'config_id': self.current_config_id}
def rotate_particles(self, angle):
pos = pyflex.get_positions().reshape(-1, 4)
center = np.mean(pos, axis=0)
pos -= center
new_pos = pos.copy()
new_pos[:, 0] = (np.cos(angle) * pos[:, 0] - np.sin(angle) * pos[:, 2])
new_pos[:, 2] = (np.sin(angle) * pos[:, 0] + np.cos(angle) * pos[:, 2])
new_pos += center
pyflex.set_positions(new_pos)
def _set_to_vertical(self, x_low, height_low):
curr_pos = pyflex.get_positions().reshape((-1, 4))
vertical_pos = self._get_vertical_pos(x_low, height_low)
curr_pos[:, :3] = vertical_pos
max_height = np.max(curr_pos[:, 1])
if max_height < 0.5:
curr_pos[:, 1] += 0.5 - max_height
pyflex.set_positions(curr_pos)
pyflex.step()
def center_object():
"""
Center the object to be at the origin
NOTE: call a pyflex.set_positions and then pyflex.step
"""
pos = pyflex.get_positions().reshape(-1, 4)
pos[:, [0, 2]] -= np.mean(pos[:, [0, 2]], axis=0, keepdims=True)
pyflex.set_positions(pos.flatten())
pyflex.step()
def get_state(self):
'''
get the postion, velocity of flex particles, and postions of flex shapes.
'''
particle_pos = pyflex.get_positions()
particle_vel = pyflex.get_velocities()
shape_position = pyflex.get_shape_states()
return {'particle_pos': particle_pos, 'particle_vel': particle_vel, 'shape_pos': shape_position,
'box_x': self.box_x, 'box_states': self.box_states, 'box_params': self.box_params, 'config_id': self.current_config_id}
def _reset(self):
""" Right now only use one initial state. Need to make sure _reset always give the same result. Otherwise CEM will fail."""
if hasattr(self, 'action_tool'):
particle_pos = pyflex.get_positions().reshape(-1, 4)
p1, p2, p3, p4 = self._get_key_point_idx()
key_point_pos = particle_pos[(
p1, p2), :3] # Was changed from from p1, p4.
middle_point = np.mean(key_point_pos, axis=0)
self.action_tool.reset([middle_point[0], 0.1, middle_point[2]])
# picker_low = self.action_tool.picker_low
# picker_high = self.action_tool.picker_high
# offset_x = self.action_tool._get_pos()[0][0][0] - picker_low[0] - 0.3
# picker_low[0] += offset_x
# picker_high[0] += offset_x
# picker_high[0] += 1.0
# self.action_tool.update_picker_boundary(picker_low, picker_high)
config = self.get_current_config()
num_particles = np.prod(config['ClothSize'], dtype=int)
particle_grid_idx = np.array(list(range(num_particles))).reshape(
config['ClothSize'][1],
config['ClothSize'][0]) # Reversed index here
cloth_dimx = config['ClothSize'][0]
x_split = cloth_dimx // 2
self.fold_group_a = particle_grid_idx[:, :x_split].flatten()
self.fold_group_b = np.flip(particle_grid_idx,
axis=1)[:, :x_split].flatten()
colors = np.zeros(num_particles)
colors[self.fold_group_a] = 1
# self.set_colors(colors) # TODO the phase actually changes the cloth dynamics so we do not change them for now. Maybe delete this later.
pyflex.step()
self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3]
pos_a = self.init_pos[self.fold_group_a, :]
pos_b = self.init_pos[self.fold_group_b, :]
self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1))
self.performance_init = None
info = self._get_info()
self.performance_init = info['performance']
return self._get_obs()
def _reset(self):
""" Right now only use one initial state"""
self.prev_dist = self._get_current_dist(pyflex.get_positions())
if hasattr(self, 'action_tool'):
particle_pos = pyflex.get_positions().reshape(-1, 4)
drop_point_pos = particle_pos[self._get_drop_point_idx(), :3]
middle_point = np.mean(drop_point_pos, axis=0)
self.action_tool.reset(
middle_point) # middle point is not really useful
picker_radius = self.action_tool.picker_radius
self.action_tool.update_picker_boundary([-0.3, 0.5, -0.5],
[0.5, 2, 0.5])
self.action_tool.set_picker_pos(picker_pos=drop_point_pos +
np.array([0., picker_radius, 0.]))
# self.action_tool.visualize_picker_boundary()
self.performance_init = None
info = self._get_info()
self.performance_init = info['performance']
return self._get_obs()
def _get_info(self):
particle_pos = pyflex.get_positions()
curr_dist = self._get_current_dist(particle_pos)
performance = -curr_dist
performance_init = performance if self.performance_init is None else self.performance_init # Use the original performance
return {
'performance':
performance,
'normalized_performance':
(performance - performance_init) / (0. - performance_init)
}
def _get_obs(self):
if self.observation_mode == 'cam_rgb':
return self.get_image(self.camera_height, self.camera_width)
if self.observation_mode == 'point_cloud':
particle_pos = np.array(pyflex.get_positions()).reshape(
[-1, 4])[:, :3].flatten()
pos = np.zeros(shape=self.particle_obs_dim, dtype=np.float)
pos[:len(particle_pos)] = particle_pos
elif self.observation_mode == 'key_point':
particle_pos = np.array(pyflex.get_positions()).reshape([-1,
4])[:, :3]
keypoint_pos = particle_pos[self.key_point_indices, :3]
pos = keypoint_pos.flatten()
# more_info = np.array([self.rope_length, self._get_endpoint_distance()])
# pos = np.hstack([more_info, pos])
# print("in _get_obs, pos is: ", pos)
if self.action_mode in ['sphere', 'picker']:
shapes = pyflex.get_shape_states()
shapes = np.reshape(shapes, [-1, 14])
pos = np.concatenate([pos.flatten(), shapes[:, 0:3].flatten()])
return pos
def _reset(self):
""" Right now only use one initial state"""
if hasattr(self, 'action_tool'):
particle_pos = pyflex.get_positions().reshape(-1, 4)
drop_point_pos = particle_pos[self._get_drop_point_idx(), :3]
middle_point = np.mean(drop_point_pos, axis=0)
self.action_tool.reset(
middle_point) # middle point is not really useful
picker_radius = self.action_tool.picker_radius
self.action_tool.update_picker_boundary([-0.3, 0.5, -0.5],
[0.5, 2, 0.5])
self.action_tool.set_picker_pos(picker_pos=drop_point_pos +
np.array([0., picker_radius, 0.]))
config = self.get_current_config()
num_particles = np.prod(config['ClothSize'], dtype=int)
particle_grid_idx = np.array(list(range(num_particles))).reshape(
config['ClothSize'][1],
config['ClothSize'][0]) # Reversed index here
cloth_dimx = config['ClothSize'][0]
x_split = cloth_dimx // 2
self.fold_group_a = particle_grid_idx[:, :x_split].flatten()
self.fold_group_b = np.flip(particle_grid_idx,
axis=1)[:, :x_split].flatten()
colors = np.zeros(num_particles)
colors[self.fold_group_a] = 1
pyflex.step()
self.init_pos = pyflex.get_positions().reshape((-1, 4))[:, :3]
pos_a = self.init_pos[self.fold_group_a, :]
pos_b = self.init_pos[self.fold_group_b, :]
self.prev_dist = np.mean(np.linalg.norm(pos_a - pos_b, axis=1))
self.performance_init = None
info = self._get_info()
self.performance_init = info['performance']
return self._get_obs()
def compute_reward(self, action=None, obs=None, set_prev_reward=False):
"""
The particles are splitted into two groups. The reward will be the minus average eculidean distance between each
particle in group a and the crresponding particle in group b
:param pos: nx4 matrix (x, y, z, inv_mass)
"""
pos = pyflex.get_positions()
pos = pos.reshape((-1, 4))[:, :3]
pos_group_a = pos[self.fold_group_a]
pos_group_b = pos[self.fold_group_b]
pos_group_b_init = self.flat_pos[self.fold_group_b]
curr_dist = np.mean(np.linalg.norm(pos_group_a - pos_group_b, axis=1)) + 1.2 * np.linalg.norm(np.mean(pos_group_b - pos_group_b_init, axis=1))
reward = -curr_dist
return reward
def _get_info(self):
# Duplicate of the compute reward function!
particle_pos = pyflex.get_positions()
curr_covered_area = self._get_current_covered_area(particle_pos)
init_covered_area = curr_covered_area if self.init_covered_area is None else self.init_covered_area
max_covered_area = self.get_current_config()['flatten_area']
info = {
'performance':
curr_covered_area,
'normalized_performance': (curr_covered_area - init_covered_area) /
(max_covered_area - init_covered_area),
}
if 'qpg' in self.action_mode:
info['total_steps'] = self.action_tool.total_steps
return info
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
