def set_scene(self, config, states=None):
self.line_box_x = self.line_box_y = None
if states is None:
super().set_scene(
config=config, states=states, create_only=False
) # this adds the water, the controlled cup, and the target cup.
else:
super().set_scene(
config=config, states=states, create_only=True
) # this adds the water, the controlled cup, and the target cup.
# needs to add an indicator box for how much water we want to pour into the target cup
# create an line on the left wall of the target cup to indicate how much water we want to pour into it.
halfEdge = np.array([
0.005 / 2., 0.005 / 2.,
(self.poured_glass_dis_z - 2 * self.poured_border) / 2.
])
center = np.array([0., 0., 0.])
quat = quatFromAxisAngle([0, 0, -1.], 0.)
print("pourwater amount add trigger box")
pyflex.add_box(
halfEdge, center, quat, 1
) # set trigger to be true to create a different color for the indicator box.
# exit()
max_water_height = self._get_current_water_height() - self.border / 2
controlled_size = (self.glass_dis_x - 2 * self.border) * (
self.glass_dis_z - 2 * self.border)
target_size = (self.poured_glass_dis_x - 2 * self.poured_border) * (
self.poured_glass_dis_z - 2 * self.poured_border)
estimated_target_water_height = max_water_height * config[
'target_amount'] / (target_size / controlled_size)
print("max_water_height in controlled cup: ", max_water_height)
print("target amount: ", config['target_amount'])
print("target size / controlled size: ", target_size / controlled_size)
print("estimated target height: ", estimated_target_water_height)
self.line_box_x = self.x_center + self.glass_distance - self.poured_glass_dis_x / 2 + self.poured_border
self.line_box_y = self.poured_border * 0.5 + estimated_target_water_height
print("line_box_y: ", self.line_box_y)
if states is None:
self.set_shape_states(self.glass_states, self.poured_glass_states)
else:
self.set_state(states)
def create_glass(self, glass_dis_x, glass_dis_z, height, border):
"""
the glass is a box, with each wall of it being a very thin box in Flex.
each wall of the real box is represented by a box object in Flex with really small thickness (determined by the param border)
dis_x: the length of the glass
dis_z: the width of the glass
height: the height of the glass.
border: the thickness of the glass wall.
the halfEdge determines the center point of each wall.
Note: this is merely setting the length of each dimension of the wall, but not the actual position of them.
That's why left and right walls have exactly the same params, and so do front and back walls.
"""
center = np.array([0., 0., 0.])
quat = quatFromAxisAngle([0, 0, -1.], 0.)
boxes = []
# floor
halfEdge = np.array([
glass_dis_x / 2. + border, border / 2., glass_dis_z / 2. + border
])
boxes.append([halfEdge, center, quat])
# left wall
halfEdge = np.array(
[border / 2., (height) / 2., glass_dis_z / 2. + border])
boxes.append([halfEdge, center, quat])
# right wall
boxes.append([halfEdge, center, quat])
# back wall
halfEdge = np.array([(glass_dis_x) / 2., (height) / 2., border / 2.])
boxes.append([halfEdge, center, quat])
# front wall
boxes.append([halfEdge, center, quat])
for i in range(len(boxes)):
halfEdge = boxes[i][0]
center = boxes[i][1]
quat = boxes[i][2]
pyflex.add_box(halfEdge, center, quat)
return boxes
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)
def create_box(self, box_dis_x, box_dis_z, height):
"""
create a box.
dis_x: the length of the box
dis_z: the width of the box
height: the height of the box.
the halfEdge determines the center point of each wall.
"""
center = np.array([0., 0., 0.])
quat = quatFromAxisAngle([0, 0, -1.], 0.)
boxes = []
# a single box
halfEdge = np.array([box_dis_x / 2., height / 2., box_dis_z / 2.])
boxes.append([halfEdge, center, quat])
for i in range(len(boxes)):
halfEdge = boxes[i][0]
center = boxes[i][1]
quat = boxes[i][2]
pyflex.add_box(halfEdge, center, quat)
return boxes
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])
print("scene_params", scene_params)
pyflex.set_scene(6, scene_params, 0)
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)
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 visualize_picker_boundary(self):
halfEdge = np.array(self.picker_high - self.picker_low) / 2.
center = np.array(self.picker_high + self.picker_low) / 2.
quat = np.array([1., 0., 0., 0.])
pyflex.add_box(halfEdge, center, quat)