def _step(self, action):
'''
action: np.ndarray of dim 1x3, (x, y, theta). (x, y) specifies the floor center coordinate, and theta
specifies the rotation.
'''
# make action as increasement, clip its range
move = action[:2]
rotate = action[2]
move = np.clip(move,
a_min=self.action_space.low[0],
a_max=self.action_space.high[0])
rotate = np.clip(rotate,
a_min=self.action_space.low[2],
a_max=self.action_space.high[2])
dx, dy, dtheta = move[0], move[1], rotate
x, y, theta = self.glass_x + dx, self.glass_y + dy, self.glass_rotation + dtheta
# check if the movement of the pouring glass collide with the poured glass.
# the action only take effects if there is no collision
new_states = self.rotate_glass(self.glass_states, x, y, theta)
if not self.judge_glass_collide(
new_states, theta) and self.above_floor(new_states, theta):
self.glass_states = new_states
self.glass_x, self.glass_y, self.glass_rotation = x, y, theta
else: # invalid move, old state becomes the same as the current state
self.glass_states[:, 3:6] = self.glass_states[:, :3].copy()
self.glass_states[:, 10:] = self.glass_states[:, 6:10].copy()
# pyflex takes a step to update the glass and the water fluid
self.set_shape_states(self.glass_states, self.poured_glass_states)
pyflex.step(render=True)
self.inner_step += 1
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 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 _step(self, action):
if self.action_mode.startswith('picker'):
self.action_tool.step(action)
pyflex.step()
else:
raise NotImplementedError
return
def _step(self, action):
self.action_tool.step(action)
if self.action_mode in ['sawyer', 'franka']:
pyflex.step(self.action_tool.next_action)
else:
pyflex.step()
return
def step(self, action):
"""
action: Array of pick_num x 4. For each picker, the action should be [x, y, z, pick/drop]. The picker will then first pick/drop, and keep
the pick/drop state while moving towards x, y, x.
"""
total_steps = 0
action = action.reshape(-1, 4)
curr_pos = np.array(pyflex.get_shape_states()).reshape(-1, 14)[:, :3]
end_pos = np.vstack([
self._apply_picker_boundary(picker_pos)
for picker_pos in action[:, :3]
])
dist = np.linalg.norm(curr_pos - end_pos, axis=1)
num_step = np.max(np.ceil(dist / self.delta_move))
if num_step < 0.1:
return
delta = (end_pos - curr_pos) / num_step
norm_delta = np.linalg.norm(delta)
for i in range(
int(min(num_step, 300))
): # The maximum number of steps allowed for one pick and place
curr_pos = np.array(pyflex.get_shape_states()).reshape(-1,
14)[:, :3]
dist = np.linalg.norm(end_pos - curr_pos, axis=1)
if np.alltrue(dist < norm_delta):
delta = end_pos - curr_pos
super().step(np.hstack([delta, action[:, 3].reshape(-1, 1)]))
pyflex.step()
total_steps += 1
if self.env is not None and self.env.recording:
self.env.video_frames.append(self.env.render(mode='rgb_array'))
if np.alltrue(dist < self.delta_move):
break
return total_steps
def _reset(self):
super()._reset()
self.performance_init = None
info = self._get_info()
self.performance_init = info['performance']
pyflex.step()
return self._get_obs()
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 _step(self, action):
#if self.action_mode == 'pickerTEO':
# action = self.action_tool.step_teo(action)
self.action_tool.step(action)
if self.action_mode in ['sawyer', 'franka']:
print(self.action_tool.next_action)
pyflex.step(self.action_tool.next_action)
else:
pyflex.step()
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 _reset(self):
'''
reset to environment to the initial state.
return the initial observation.
'''
self.inner_step = 0
self.performance_init = None
info = self._get_info()
self.performance_init = info['performance']
pyflex.step(render=True)
return self._get_obs()
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 _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 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 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 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 _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 _step(self, action):
'''
action: np.ndarray of dim 1x1, dx, which specifies how much to move on the x-axis.
'''
# make action as increasement, clip its range
dx = action[0]
dx = np.clip(dx, a_min=self.action_space.low[0], a_max=self.action_space.high[0])
x = self.box_x + dx
# move the box
new_states = self.move_box(self.box_states, x)
self.box_states = new_states
self.box_x = x
self.box_x = np.clip(self.box_x, a_min=self.min_x, a_max=self.max_x)
# pyflex takes a step to update the box and the water fluid
pyflex.set_shape_states(self.box_states)
pyflex.step()
self.inner_step += 1
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 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 _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 _set_to_folded(self):
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
fold_group_a = particle_grid_idx[:, :x_split].flatten()
fold_group_b = np.flip(particle_grid_idx,
axis=1)[:, :x_split].flatten()
curr_pos = pyflex.get_positions().reshape((-1, 4))
curr_pos[fold_group_a, :] = curr_pos[fold_group_b, :].copy()
curr_pos[
fold_group_a,
1] += 0.05 # group a particle position made tcurr_pos[self.fold_group_b, 1] + 0.05e at top of group b position.
pyflex.set_positions(curr_pos)
for i in range(10):
pyflex.step()
return self._get_info()['performance']
def step(self, action):
""" Action is in 5D: (u,v) the start of the pick in image coordinate; (dx, dy, dz): the relative position of the place w.r.t. the pick"""
u, v = action[:2]
u = ((u + 1.) * 0.5) * self.image_size[0]
v = ((v + 1.) * 0.5) * self.image_size[1]
x, y, z = self._get_world_coor_from_image(u, v)
y += 0.01
dx, dy, dz = action[2:]
st_high = np.array([x, 0.2, z, 0])
st = np.array([x, y, z, 0])
en = st + np.array([dx, dy, dz, 1])
# print('st:', st)
if self.full:
self.total_steps += super().step(st_high)
self.total_steps += super().step(st)
self.total_steps += super().step(en)
en[3] = 0 # Drop cloth
# Unpick all particles
_, particle_pos = self._get_pos()
new_particle_pos = particle_pos.copy()
for i in range(self.num_picker):
if self.picked_particles[i] is not None:
new_particle_pos[
self.picked_particles[i], 3] = self.particle_inv_mass[
self.picked_particles[i]] # Revert the mass
self.picked_particles[i] = None
pyflex.set_positions(new_particle_pos)
for i in range(20):
pyflex.step()
if self.env is not None and self.env.recording:
self.env.video_frames.append(
self.env.render(mode='rgb_array'))
self.total_steps += 20
else:
raise NotImplementedError
return self.total_steps
def set_scene(self, config, states=None, create_only=False):
'''
Construct the pouring water scence.
'''
# create fluid
super().set_scene(
config
) # do not sample fluid parameters, as it's very likely to generate very strange fluid
# compute glass params
if states is None:
self.set_pouring_glass_params(config["glass"])
self.set_poured_glass_params(config["glass"])
else:
glass_params = states['glass_params']
self.border = glass_params['border']
self.height = glass_params['height']
self.glass_dis_x = glass_params['glass_dis_x']
self.glass_dis_z = glass_params['glass_dis_z']
self.glass_distance = glass_params['glass_distance']
self.poured_border = glass_params['poured_border']
self.poured_height = glass_params['poured_height']
self.poured_glass_dis_x = glass_params['poured_glass_dis_x']
self.poured_glass_dis_z = glass_params['poured_glass_dis_z']
self.glass_params = glass_params
# create pouring glass & poured glass
self.create_glass(self.glass_dis_x, self.glass_dis_z, self.height,
self.border)
self.create_glass(self.poured_glass_dis_x, self.poured_glass_dis_z,
self.poured_height, self.poured_border)
# move pouring glass to be at ground
self.glass_states = self.init_glass_state(self.x_center, 0,
self.glass_dis_x,
self.glass_dis_z,
self.height, self.border)
# move poured glass to be at ground
self.poured_glass_states = self.init_glass_state(
self.x_center + self.glass_distance, 0, self.poured_glass_dis_x,
self.poured_glass_dis_z, self.poured_height, self.poured_border)
self.set_shape_states(self.glass_states, self.poured_glass_states)
# record glass floor center x, y, and rotation
self.glass_x = self.x_center
if self.action_mode == 'rotation_bottom':
self.glass_y = 0
elif self.action_mode == 'rotation_top':
self.glass_y = 0.5 * self.border + self.height
self.glass_rotation = 0
# only create the glass and water, without setting their states
# this is only used in the pourwater amount env.
if create_only:
return
# no cached init states passed in
if states is None:
fluid_pos = np.ones((self.particle_num, self.dim_position))
# move water all inside the glass
fluid_radius = self.fluid_params['radius'] * self.fluid_params[
'rest_dis_coef']
fluid_dis = np.array(
[1.0 * fluid_radius, fluid_radius * 0.5, 1.0 * fluid_radius])
lower_x = self.glass_params['glass_x_center'] - self.glass_params[
'glass_dis_x'] / 2. + self.glass_params['border']
lower_z = -self.glass_params[
'glass_dis_z'] / 2 + self.glass_params['border']
lower_y = self.glass_params['border']
if self.action_mode in ['sawyer', 'franka']:
lower_y += 0.56 # NOTE: robotics table
lower = np.array([lower_x, lower_y, lower_z])
cnt = 0
rx = int(self.fluid_params['dim_x'] * 1)
ry = int(self.fluid_params['dim_y'] * 1)
rz = int(self.fluid_params['dim_z'] / 1)
for x in range(rx):
for y in range(ry):
for z in range(rz):
fluid_pos[cnt][:3] = lower + np.array(
[x, y, z]) * fluid_dis # + np.random.rand() * 0.01
cnt += 1
pyflex.set_positions(fluid_pos)
print("stablize water!")
for _ in range(100):
pyflex.step()
state_dic = self.get_state()
water_state = state_dic['particle_pos'].reshape(
(-1, self.dim_position))
in_glass = self.in_glass(water_state, self.glass_states,
self.border, self.height)
not_in_glass = 1 - in_glass
not_total_num = np.sum(not_in_glass)
while not_total_num > 0:
max_height_now = np.max(water_state[:, 1])
fluid_dis = np.array(
[1.0 * fluid_radius, fluid_radius * 1, 1.0 * fluid_radius])
lower_x = self.glass_params[
'glass_x_center'] - self.glass_params['glass_dis_x'] / 4
lower_z = -self.glass_params['glass_dis_z'] / 4
lower_y = max_height_now
lower = np.array([lower_x, lower_y, lower_z])
cnt = 0
dim_x = config['fluid']['dim_x']
dim_z = config['fluid']['dim_z']
for w_idx in range(len(water_state)):
if not in_glass[w_idx]:
water_state[w_idx][:3] = lower + fluid_dis * np.array([
cnt % dim_x, cnt // (dim_x * dim_z),
(cnt // dim_x) % dim_z
])
cnt += 1
pyflex.set_positions(water_state)
for _ in range(40):
pyflex.step()
state_dic = self.get_state()
water_state = state_dic['particle_pos'].reshape(
(-1, self.dim_position))
in_glass = self.in_glass(water_state, self.glass_states,
self.border, self.height)
not_in_glass = 1 - in_glass
not_total_num = np.sum(not_in_glass)
for _ in range(30):
pyflex.step()
else: # set to passed-in cached init states
self.set_state(states)
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")
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
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 gen_PyFleX(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']
shape_state_dim, dt = info['shape_state_dim'], info['dt']
gen_vision = info['gen_vision']
vision_dir, vis_width, vis_height = info['vision_dir'], info['vis_width'], info['vis_height']
np.random.seed(round(time.time() * 1000 + thread_idx) % 2 ** 32)
# positions
stats = [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 == 'RigidFall':
g_low, g_high = info['physics_param_range']
gravity = rand_float(g_low, g_high)
print("Generated RigidFall rollout {} with gravity {} from range {} ~ {}".format(
i, gravity, g_low, g_high))
n_instance = 3
draw_mesh = 1
scene_params = np.zeros(n_instance * 3 + 3)
scene_params[0] = n_instance
scene_params[1] = gravity
scene_params[-1] = draw_mesh
low_bound = 0.09
for j in range(n_instance):
x = rand_float(0., 0.1)
y = rand_float(low_bound, low_bound + 0.01)
z = rand_float(0., 0.1)
scene_params[j * 3 + 2] = x
scene_params[j * 3 + 3] = y
scene_params[j * 3 + 4] = z
low_bound += 0.21
pyflex.set_scene(env_idx, scene_params, thread_idx)
pyflex.set_camPos(np.array([0.2, 0.875, 2.0]))
n_particles = pyflex.get_n_particles()
n_shapes = 1 # the floor
positions = np.zeros((time_step, n_particles + n_shapes, 3), dtype=np.float32)
shape_quats = np.zeros((time_step, n_shapes, 4), dtype=np.float32)
for j in range(time_step):
positions[j, :n_particles] = pyflex.get_positions().reshape(-1, 4)[:, :3]
ref_positions = positions[0]
for k in range(n_instance):
XX = ref_positions[64*k:64*(k+1)]
YY = positions[j, 64*k:64*(k+1)]
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*k:64*(k+1)] = YY_fitted
if gen_vision:
pyflex.step(capture=True, path=os.path.join(rollout_dir, str(j) + '.tga'))
else:
pyflex.step()
data = [positions[j], shape_quats[j], scene_params]
store_data(data_names, data, os.path.join(rollout_dir, str(j) + '.h5'))
if gen_vision:
images = np.zeros((time_step, vis_height, vis_width, 3), dtype=np.uint8)
for j in range(time_step):
img_path = os.path.join(rollout_dir, str(j) + '.tga')
img = scipy.misc.imread(img_path)[:, :, :3][:, :, ::-1]
img = cv2.resize(img, (vis_width, vis_height), interpolation=cv2.INTER_AREA)
images[j] = img
os.system('rm ' + img_path)
store_data(['positions', 'images', 'scene_params'], [positions, images, scene_params],
os.path.join(vision_dir, str(rollout_idx) + '.h5'))
elif env == 'MassRope':
s_low, s_high = info['physics_param_range']
stiffness = rand_float(s_low, s_high)
print("Generated MassRope rollout {} with gravity {} from range {} ~ {}".format(
i, stiffness, s_low, s_high))
x = 0.
y = 1.0
z = 0.
length = 0.7
draw_mesh = 1.
scene_params = np.array([x, y, z, length, stiffness, draw_mesh])
pyflex.set_scene(env_idx, scene_params, 0)
pyflex.set_camPos(np.array([0.13, 2.0, 3.2]))
action = np.zeros(3)
# the last particle is the pin, regarded as shape
n_particles = pyflex.get_n_particles() - 1
n_shapes = 1 # the mass at the top of the rope
positions = np.zeros((time_step + 1, n_particles + n_shapes, 3), dtype=np.float32)
shape_quats = np.zeros((time_step + 1, n_shapes, 4), dtype=np.float32)
action = np.zeros(3)
for j in range(time_step + 1):
positions[j] = pyflex.get_positions().reshape(-1, 4)[:, :3]
if j >= 1:
# append the action (position of the pin) to the previous time step
positions[j - 1, -1, :] = positions[j, -1, :]
ref_positions = positions[0]
# apply rigid projection to the rigid object
# cube: [0, 81)
# rope: [81, 95)
# pin: [95, 96)
XX = ref_positions[:81]
YY = positions[j, :81]
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
positions[j, :81] = YY_fitted
scale = 0.1
action[0] += rand_float(-scale, scale) - positions[j, -1, 0] * 0.1
action[2] += rand_float(-scale, scale) - positions[j, -1, 2] * 0.1
if gen_vision:
pyflex.step(action * dt, capture=True, path=os.path.join(rollout_dir, str(j) + '.tga'))
else:
pyflex.step(action * dt)
if j >= 1:
data = [positions[j - 1], shape_quats[j - 1], scene_params]
store_data(data_names, data, os.path.join(rollout_dir, str(j - 1) + '.h5'))
if gen_vision:
images = np.zeros((time_step, vis_height, vis_width, 3), dtype=np.uint8)
for j in range(time_step):
img_path = os.path.join(rollout_dir, str(j) + '.tga')
img = scipy.misc.imread(img_path)[:, :, :3][:, :, ::-1]
img = cv2.resize(img, (vis_width, vis_height), interpolation=cv2.INTER_AREA)
images[j] = img
os.system('rm ' + img_path)
store_data(['positions', 'images', 'scene_params'], [positions, images, scene_params],
os.path.join(vision_dir, str(rollout_idx) + '.h5'))
else:
raise AssertionError("Unsupported env")
# change dtype for more accurate stat calculation
# only normalize positions
datas = [positions[:time_step].astype(np.float64)]
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 reset(self, state):
for i in range(100):
pyflex.step()
