def move_box(self, prev_states, x):
'''
given the previous states of the box, move it in 1D along x-axis.
state:
0-3: current (x, y, z) coordinate of the center point
3-6: previous (x, y, z) coordinate of the center point
6-10: current quat
10-14: previous quat
'''
quat_curr = quatFromAxisAngle([0, 0, -1.], 0.)
states = np.zeros((1, self.dim_shape_state))
for i in range(1):
states[i][3:6] = prev_states[i][:3]
states[i][10:] = prev_states[i][6:10]
x_center = x
y = self.height / 2.
# floor: center position does not change
states[0, :3] = np.array([x_center, y, 0.])
states[:, 6:10] = quat_curr
return states
def init_glass_state(self, x, y, glass_dis_x, glass_dis_z, height, border):
'''
set the initial state of the glass.
'''
dis_x, dis_z = glass_dis_x, glass_dis_z
x_center, y_curr, y_last = x, y, 0.
quat = quatFromAxisAngle([0, 0, -1.], 0.)
# states of 5 walls
states = np.zeros((5, self.dim_shape_state))
# floor
states[0, :3] = np.array([x_center, y_curr, 0.])
states[0, 3:6] = np.array([x_center, y_last, 0.])
# left wall
states[1, :3] = np.array([x_center - (dis_x + border) / 2., (height + border) / 2. + y_curr, 0.])
states[1, 3:6] = np.array([x_center - (dis_x + border) / 2., (height + border) / 2. + y_last, 0.])
# right wall
states[2, :3] = np.array([x_center + (dis_x + border) / 2., (height + border) / 2. + y_curr, 0.])
states[2, 3:6] = np.array([x_center + (dis_x + border) / 2., (height + border) / 2. + y_last, 0.])
# back wall
states[3, :3] = np.array([x_center, (height + border) / 2. + y_curr, -(dis_z + border) / 2.])
states[3, 3:6] = np.array([x_center, (height + border) / 2. + y_last, -(dis_z + border) / 2.])
# front wall
states[4, :3] = np.array([x_center, (height + border) / 2. + y_curr, (dis_z + border) / 2.])
states[4, 3:6] = np.array([x_center, (height + border) / 2. + y_last, (dis_z + border) / 2.])
states[:, 6:10] = quat
states[:, 10:] = quat
return states
def move_glass(self, prev_states, x):
'''
given the previous states of the glass, move it in 1D along x-axis.
update the states of the 5 boxes that form the box: floor, left/right wall, back/front wall.
state:
0-3: current (x, y, z) coordinate of the center point
3-6: previous (x, y, z) coordinate of the center point
6-10: current quat
10-14: previous quat
'''
dis_x, dis_z = self.glass_dis_x, self.glass_dis_z
quat_curr = quatFromAxisAngle([0, 0, -1.], 0.)
border = self.border
# states of 5 walls
states = np.zeros((5, self.dim_shape_state))
for i in range(5):
states[i][3:6] = prev_states[i][:3]
states[i][10:] = prev_states[i][6:10]
x_center = x
y = 0
# floor: center position does not change
states[0, :3] = np.array([x_center, y, 0.])
# left wall: center must move right and move down.
relative_coord = np.array(
[-(dis_x + border) / 2., (self.height + border) / 2., 0.])
states[1, :3] = states[0, :3] + relative_coord
# right wall
relative_coord = np.array([(dis_x + border) / 2.,
(self.height + border) / 2., 0.])
states[2, :3] = states[0, :3] + relative_coord
# back wall
relative_coord = np.array(
[0, (self.height + border) / 2., -(dis_z + border) / 2.])
states[3, :3] = states[0, :3] + relative_coord
# front wall
relative_coord = np.array(
[0, (self.height + border) / 2., (dis_z + border) / 2.])
states[4, :3] = states[0, :3] + relative_coord
states[:, 6:10] = quat_curr
return 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_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 init_box_state(self, x, y, box_dis_x, box_dis_z, height):
'''
set the initial state of the box.
'''
dis_x, dis_z = box_dis_x, box_dis_z
x_center, y_curr, y_last = x, y, 0.
if self.action_mode in ['sawyer', 'franka']:
y_curr = y_last = 0.56 # NOTE: robotics table
quat = quatFromAxisAngle([0, 0, -1.], 0.)
# states of a single box
states = np.zeros((1, self.dim_shape_state))
# a single box
states[0, :3] = np.array([x_center, height / 2., 0.])
states[0, 3:6] = np.array([x_center, height / 2., 0.])
states[:, 6:10] = quat
states[:, 10:] = quat
return states
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 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
def rotate_glass(self, prev_states, x, y, theta):
'''
given the previous states of the glass, rotate it with angle theta.
update the states of the 5 boxes that form the box: floor, left/right wall, back/front wall.
rotate the glass, where the center point is the center of the floor or the top.
state:
0-3: current (x, y, z) coordinate of the center point
3-6: previous (x, y, z) coordinate of the center point
6-10: current quat
10-14: previous quat
'''
dis_x, dis_z = self.glass_dis_x, self.glass_dis_z
quat_curr = quatFromAxisAngle([0, 0, -1.], theta)
border = self.border
# states of 5 walls
states = np.zeros((5, self.dim_shape_state))
for i in range(5):
states[i][3:6] = prev_states[i][:3]
states[i][10:] = prev_states[i][6:10]
x_center = x
# rotation center is the floor center
rotate_center = np.array([x_center, y, 0.])
if self.action_mode == 'rotation_bottom':
# floor: center position does not change
states[0, :3] = np.array([x_center, y, 0.])
# left wall: center must move right and move down.
relative_coord = np.array(
[-(dis_x + border) / 2., (self.height) / 2., 0.])
states[1, :3] = rotate_rigid_object(center=rotate_center,
axis=np.array([0, 0, -1]),
angle=theta,
relative=relative_coord)
# right wall
relative_coord = np.array([(dis_x + border) / 2.,
(self.height) / 2., 0.])
states[2, :3] = rotate_rigid_object(center=rotate_center,
axis=np.array([0, 0, -1]),
angle=theta,
relative=relative_coord)
# back wall
relative_coord = np.array(
[0, (self.height) / 2., -(dis_z + border) / 2.])
states[3, :3] = rotate_rigid_object(center=rotate_center,
axis=np.array([0, 0, -1]),
angle=theta,
relative=relative_coord)
# front wall
relative_coord = np.array(
[0, (self.height) / 2., (dis_z + border) / 2.])
states[4, :3] = rotate_rigid_object(center=rotate_center,
axis=np.array([0, 0, -1]),
angle=theta,
relative=relative_coord)
elif self.action_mode == 'rotation_top':
# floor
relative_coord = np.array([0, -self.height, 0.])
states[0, :3] = rotate_rigid_object(center=rotate_center,
axis=np.array([0, 0, -1]),
angle=theta,
relative=relative_coord)
# left wall
relative_coord = np.array(
[-(dis_x + border) / 2., -self.height / 2., 0.])
states[1, :3] = rotate_rigid_object(center=rotate_center,
axis=np.array([0, 0, -1]),
angle=theta,
relative=relative_coord)
# right wall
relative_coord = np.array([(dis_x + border) / 2.,
-self.height / 2., 0.])
states[2, :3] = rotate_rigid_object(center=rotate_center,
axis=np.array([0, 0, -1]),
angle=theta,
relative=relative_coord)
# back wall
relative_coord = np.array(
[0, -self.height / 2., -(dis_z + border) / 2.])
states[3, :3] = rotate_rigid_object(center=rotate_center,
axis=np.array([0, 0, -1]),
angle=theta,
relative=relative_coord)
# front wall
relative_coord = np.array([0, -self.height / 2., (dis_z) / 2.])
states[4, :3] = rotate_rigid_object(center=rotate_center,
axis=np.array([0, 0, -1]),
angle=theta,
relative=relative_coord)
states[:, 6:10] = quat_curr
return states