def evaluate_mpc(model_dir,
config_planner_mpc=None,
save_dir=None,
planner_type=None,
env_config=None,
strict=True,
generate_initial_condition_func=None,
):
assert save_dir is not None
assert planner_type is not None
assert env_config is not None
assert generate_initial_condition_func is not None
model_data = model_builder.load_dynamics_model_from_folder(model_dir, strict=strict)
model_dy = model_data['model_dy']
model_dy = model_dy.eval()
model_dy = model_dy.cuda()
model_config = model_dy.config
# create the environment
env = DrakePusherSliderEnv(env_config, visualize=False)
env.reset()
observation_function = ObservationFunctionFactory.function_from_config(model_config)
action_function = ActionFunctionFactory.function_from_config(model_config)
episode = OnlineEpisodeReader()
mpc_input_builder = DynamicsModelInputBuilder(observation_function=observation_function,
action_function=action_function,
visual_observation_function=None,
episode=episode)
pts_GT = np.array(model_config['dataset']['observation_function']['GT_3D_object_points'])
K = pts_GT.shape[0] # num keypoints
eval_indices = np.arange(3 * K) # extract the keypoints
def goal_func(obs_local):
"""
Helper function for getting the goal state from an observation
"""
return observation_function(obs_local)[eval_indices]
# make a planner config, same as model config but with mpc and eval sections
# replaced
planner_config = copy.copy(model_config)
if config_planner_mpc is not None:
planner_config['mpc'] = config_planner_mpc['mpc']
planner_config['eval'] = config_planner_mpc['eval']
planner = None
if planner_type == "random_shooting":
planner = RandomShootingPlanner(planner_config)
elif planner_type == "mppi":
planner = PlannerMPPI(planner_config)
else:
raise ValueError("unknown planner type: %s" %(planner_type))
# run a single iteration
mpc_eval_drake_pusher_slider.evaluate_mpc(model_dy=model_dy,
env=env,
episode=episode,
mpc_input_builder=mpc_input_builder,
planner=planner,
eval_indices=eval_indices,
goal_func=goal_func,
config=planner_config,
wait_for_user_input=False,
save_dir=save_dir,
model_name="test",
experiment_name="test",
generate_initial_condition_func=generate_initial_condition_func,
)
return {'save_dir': save_dir}
def evaluate_mpc_z_state(
model_dir,
config_planner_mpc=None,
save_dir=None,
planner_type=None,
env_config=None,
strict=True,
generate_initial_condition_func=None,
env_type="DrakePusherSliderEnv",
):
assert save_dir is not None
assert planner_type is not None
assert env_config is not None
assert generate_initial_condition_func is not None
model_dict = load_model(model_dir, strict=strict)
model_dy = model_dict['model_dy']['model_dy']
config = model_dict['model_dy']['config']
model_config = config
model_kp = model_dict['model_kp']['model']
config_kp = model_kp.config
camera_name = config_kp['perception']['camera_name']
# create the environment
# create the environment
env = None
if env_type == "DrakePusherSliderEnv":
env = DrakePusherSliderEnv(env_config, visualize=False)
elif env_type == "DrakeMugsEnv":
env = DrakeMugsEnv(env_config, visualize=False)
else:
raise ValueError("unknown env type: %s" % (env_type))
env.reset()
T_world_camera = env.camera_pose(camera_name)
camera_K_matrix = env.camera_K_matrix(camera_name)
mask_labels = env.get_labels_to_mask_list()
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.drake_pusher_position_3D(
config)
visual_observation_function = \
VisualObservationFunctionFactory.function_from_config(config,
camera_name=camera_name,
model_kp=model_kp,
K_matrix=camera_K_matrix,
T_world_camera=T_world_camera,
mask_labels=mask_labels)
episode = OnlineEpisodeReader()
mpc_input_builder = DynamicsModelInputBuilder(
observation_function=observation_function,
visual_observation_function=visual_observation_function,
action_function=action_function,
episode=episode)
def goal_func(obs_tmp):
state_tmp = mpc_input_builder.get_state_input_single_timestep(
{'observation': obs_tmp})['state']
# z_dict= model_dy.compute_z_state(state_tmp.unsqueeze(0))
# print("z_dict['z_object'].shape", z_dict['z_object'].shape)
return model_dy.compute_z_state(
state_tmp.unsqueeze(0))['z_object_flat']
index_dict = get_object_and_robot_state_indices(model_config)
object_indices = index_dict['object_indices']
# make a planner config, same as model config but with mpc and eval sections
# replaced
planner_config = copy.copy(model_config)
if config_planner_mpc is not None:
planner_config['mpc'] = config_planner_mpc['mpc']
planner_config['eval'] = config_planner_mpc['eval']
planner = None
if planner_type == "random_shooting":
planner = RandomShootingPlanner(planner_config)
elif planner_type == "mppi":
planner = PlannerMPPI(planner_config)
else:
raise ValueError("unknown planner type: %s" % (planner_type))
# run a single iteration
mpc_eval_drake_pusher_slider.evaluate_mpc(
model_dy=model_dy,
env=env,
episode=episode,
mpc_input_builder=mpc_input_builder,
planner=planner,
eval_indices=object_indices,
goal_func=goal_func,
config=planner_config,
wait_for_user_input=False,
save_dir=save_dir,
model_name="test",
experiment_name="test",
generate_initial_condition_func=generate_initial_condition_func)
return {'save_dir': save_dir}
def evaluate_mpc(
model_dir,
config_planner_mpc=None,
save_dir=None,
planner_type=None,
env_config=None,
strict=True,
generate_initial_condition_func=None,
):
assert save_dir is not None
assert planner_type is not None
assert env_config is not None
assert generate_initial_condition_func is not None
model_data = load_model(model_dir, strict=strict)
model_dy = model_data['model_dy']
model_config = model_dy.config
config = model_config
model_dd = model_data['model_dd']
# create the environment
env = DrakePusherSliderEnv(env_config, visualize=False)
env.reset()
camera_name = model_data['metadata']['camera_name']
# sanity check
camera_name_in_training = model_config['dataset'][
'visual_observation_function']['camera_name']
assert camera_name == camera_name_in_training, "camera_names don't match: camera_name = %s, camera_name_in_trainig = %s" % (
camera_name, camera_name_in_training)
T_world_camera = env.camera_pose(camera_name)
camera_K_matrix = env.camera_K_matrix(camera_name)
spatial_descriptor_data = model_data['spatial_descriptor_data']
ref_descriptors = torch.Tensor(
spatial_descriptor_data['spatial_descriptors']).cuda()
K = ref_descriptors.shape[0]
action_function = ActionFunctionFactory.function_from_config(config)
observation_function = ObservationFunctionFactory.drake_pusher_position_3D(
config)
visual_observation_function = \
VisualObservationFunctionFactory.function_from_config(config,
camera_name=camera_name,
model_dd=model_dd,
ref_descriptors=ref_descriptors,
K_matrix=camera_K_matrix,
T_world_camera=T_world_camera,
)
episode = OnlineEpisodeReader()
mpc_input_builder = DynamicsModelInputBuilder(
observation_function=observation_function,
visual_observation_function=visual_observation_function,
action_function=action_function,
episode=episode)
def goal_func(obs_local):
keypoints_dict = visual_observation_function(obs_local)
return keypoints_dict['tensor']
eval_indices = np.arange(3 * K) # extract the keypoints
# make a planner config, same as model config but with mpc and eval sections
# replaced
planner_config = copy.copy(model_config)
if config_planner_mpc is not None:
planner_config['mpc'] = config_planner_mpc['mpc']
planner_config['eval'] = config_planner_mpc['eval']
planner = None
if planner_type == "random_shooting":
planner = RandomShootingPlanner(planner_config)
elif planner_type == "mppi":
planner = PlannerMPPI(planner_config)
else:
raise ValueError("unknown planner type: %s" % (planner_type))
# run a single iteration
mpc_eval_drake_pusher_slider.evaluate_mpc(
model_dy=model_dy,
env=env,
episode=episode,
mpc_input_builder=mpc_input_builder,
planner=planner,
eval_indices=eval_indices,
goal_func=goal_func,
config=planner_config,
wait_for_user_input=False,
save_dir=save_dir,
model_name="test",
experiment_name="test",
generate_initial_condition_func=generate_initial_condition_func)
return {'save_dir': save_dir}