def __init__(self, observation_mode, action_mode, config=None, cached_states_path='transport_torus_init_states.pkl', **kwargs):
'''
This class implements a transport torus task.
The torus is put on a box. You need to move the box to a target location.
This is a 1D task.
observation_mode: "cam_rgb" or "full_state"
action_mode: "direct"
TODO: add more description of the task.
'''
assert observation_mode in ['cam_rgb', 'point_cloud', 'key_point']
self.observation_mode = observation_mode
self.action_mode = action_mode
self.wall_num = 5 # number of box walls. floor/left/right/front/back
self.distance_coef = 1.
self.torus_penalty_coef = 10.
self.terminal_x = 1.2
self.min_x = -0.25
self.max_x = 1.4
self.prev_reward = 0
self.reward_min = self.distance_coef * min(self.min_x - self.terminal_x, self.terminal_x - self.max_x) - \
self.torus_penalty_coef * 1
self.reward_max = 0 # reach target position, with no water spilled
self.reward_range = self.reward_max - self.reward_min
super().__init__(**kwargs)
self.get_cached_configs_and_states(cached_states_path, self.num_variations)
if observation_mode in ['point_cloud', 'key_point']:
if observation_mode == 'key_point':
obs_dim = 0
else:
max_particle_num = 13 * 13 * 13 * 4
obs_dim = max_particle_num * 3
self.particle_obs_dim = obs_dim
# z and theta of the second cup (poured_box) does not change and thus are omitted.
obs_dim += 8 # Pos (x) and shape (w, h, l) reset of the cup, torus x, torus y, if on box.
self.observation_space = Box(low=np.array([-np.inf] * obs_dim), high=np.array([np.inf] * obs_dim), dtype=np.float32)
elif observation_mode == 'cam_rgb':
self.observation_space = Box(low=-np.inf, high=np.inf, shape=(self.camera_height, self.camera_width, 3),
dtype=np.float32)
default_config = self.get_default_config()
if action_mode == 'direct':
self.action_direct_dim = 1
action_low = np.array([-0.011])
action_high = np.array([0.011])
self.action_space = Box(action_low, action_high, dtype=np.float32)
elif action_mode in ['sawyer', 'franka']:
self.action_tool = RobotBase(action_mode)
else:
raise NotImplementedError
def __init__(self,
observation_mode,
action_mode,
num_picker=2,
horizon=75,
render_mode='particle',
picker_radius=0.02,
**kwargs):
self.render_mode = render_mode
super().__init__(**kwargs)
assert observation_mode in ['point_cloud', 'cam_rgb', 'key_point']
assert action_mode in ['picker', 'sawyer', 'franka']
self.observation_mode = observation_mode
self.action_mode = action_mode
self.num_picker = num_picker
if action_mode == 'picker':
self.action_tool = Picker(num_picker,
picker_radius=picker_radius,
picker_threshold=0.005,
particle_radius=0.025,
picker_low=(-0.35, 0., -0.35),
picker_high=(0.35, 0.3, 0.35))
self.action_space = self.action_tool.action_space
elif action_mode in ['sawyer', 'franka']:
self.action_tool = RobotBase(action_mode)
if observation_mode in ['key_point', 'point_cloud']:
if observation_mode == 'key_point':
obs_dim = 10 * 3
else:
max_particles = 41
obs_dim = max_particles * 3
self.particle_obs_dim = obs_dim
if action_mode in ['picker']:
obs_dim += num_picker * 3
else:
raise NotImplementedError
self.observation_space = Box(np.array([-np.inf] * obs_dim),
np.array([np.inf] * obs_dim),
dtype=np.float32)
elif observation_mode == 'cam_rgb':
self.observation_space = Box(low=-np.inf,
high=np.inf,
shape=(self.camera_height,
self.camera_width, 3),
dtype=np.float32)
self.horizon = horizon
def __init__(self, observation_mode="cam_rgb", action_mode="picker", num_picker=2, render_mode='particle', picker_radius=0.05, particle_radius=0.00625, **kwargs):
self.render_mode = render_mode
self.action_mode = action_mode
self.cloth_particle_radius = particle_radius
super().__init__(**kwargs)
assert observation_mode in ['key_point', 'point_cloud', 'cam_rgb']
assert action_mode in ['picker', 'pickerpickplace', 'sawyer', 'franka', 'picker_qpg', 'pickerTEO']
self.observation_mode = observation_mode
if action_mode == 'picker':
self.action_tool = Picker(num_picker, picker_radius=picker_radius, particle_radius=particle_radius,
picker_low=(-0.4, 0., -0.4), picker_high=(1.0, 0.5, 0.4))
self.action_space = self.action_tool.action_space
self.picker_radius = picker_radius
elif action_mode == 'pickerpickplace':
self.action_tool = PickerPickPlace(num_picker=num_picker, particle_radius=particle_radius, env=self,
picker_low=(-0.5, 0., -0.5), picker_high=(0.5, 0.3, 0.5))
self.action_space = self.action_tool.action_space
assert self.action_repeat == 1
elif action_mode in ['sawyer', 'franka']:
self.action_tool = RobotBase(action_mode)
self.action_space = self.action_tool.action_space
elif action_mode == 'picker_qpg':
cam_pos, cam_angle = self.get_camera_params()
self.action_tool = PickerQPG((self.camera_height, self.camera_height), cam_pos, cam_angle,
num_picker=num_picker, particle_radius=particle_radius, env=self,
picker_low=(-0.3, 0., -0.3), picker_high=(0.3, 0.3, 0.3)
)
self.action_space = self.action_tool.action_space
if observation_mode in ['key_point', 'point_cloud']:
if observation_mode == 'key_point':
obs_dim = len(self._get_key_point_idx()) * 3
else:
max_particles = 120 * 120
obs_dim = max_particles * 3
self.particle_obs_dim = obs_dim
if action_mode.startswith('picker'):
obs_dim += num_picker * 3
else:
raise NotImplementedError
self.observation_space = Box(np.array([-np.inf] * obs_dim), np.array([np.inf] * obs_dim), dtype=np.float32)
elif observation_mode == 'cam_rgb':
self.observation_space = Box(low=-np.inf, high=np.inf, shape=(self.camera_height, self.camera_width, 3),
dtype=np.float32)
def __init__(self,
observation_mode,
action_mode,
cached_states_path='pass_water_init_states.pkl',
**kwargs):
'''
This class implements a pouring water task.
observation_mode: "cam_rgb" or "full_state"
action_mode: "direct"
TODO: add more description of the task.
'''
assert observation_mode in ['cam_rgb', 'point_cloud', 'key_point']
self.observation_mode = observation_mode
self.action_mode = action_mode
self.wall_num = 5 # number of glass walls. floor/left/right/front/back
self.distance_coef = 1.
self.water_penalty_coef = 10.
self.terminal_x = 1.2
self.min_x = -0.25
self.max_x = 1.4
self.prev_reward = 0
self.reward_min = self.distance_coef * min(self.min_x - self.terminal_x, self.terminal_x - self.max_x) - \
self.water_penalty_coef * 1
self.reward_max = 0 # reach target position, with no water spilled
self.reward_range = self.reward_max - self.reward_min
super().__init__(**kwargs)
self.get_cached_configs_and_states(cached_states_path,
self.num_variations)
if observation_mode in ['point_cloud', 'key_point']:
if observation_mode == 'key_point':
# z and theta of the second cup (poured_glass) does not change and thus are omitted.
obs_dim = 7 # Pos (x) and shape (w, h, l) reset of the cup, as well as the water height.
# also: frac of water in, frac of water out
else:
max_particle_num = 13 * 13 * 13 * 4
obs_dim = max_particle_num * 3
self.particle_obs_dim = obs_dim
self.observation_space = Box(low=np.array([-np.inf] * obs_dim),
high=np.array([np.inf] * obs_dim),
dtype=np.float32)
elif observation_mode == 'cam_rgb':
self.observation_space = Box(low=-np.inf,
high=np.inf,
shape=(self.camera_height,
self.camera_width, 3),
dtype=np.float32)
default_config = self.get_default_config()
# raidus = default_config['fluid']['radius'] * default_config['fluid']['rest_dis_coef']
border = default_config['glass'][
'border'] #* default_config['fluid']['rest_dis_coef']
if action_mode == 'direct':
self.action_direct_dim = 1
action_low = np.array([-border * 0.5])
action_high = np.array([border * 0.5])
self.action_space = Box(action_low, action_high, dtype=np.float32)
elif action_mode in ['sawyer', 'franka']:
self.action_tool = RobotBase(action_mode)
else:
raise NotImplementedError