def main():
colorama.init(autoreset=True)
plt.style.use("slides")
np.set_printoptions(precision=3, suppress=True, linewidth=200)
parser = argparse.ArgumentParser(formatter_class=my_formatter)
parser.add_argument("fwd_model_dir",
help="load this saved forward model file",
type=pathlib.Path,
nargs='+')
parser.add_argument("test_config",
help="json file describing the test",
type=pathlib.Path)
parser.add_argument("labeling_params",
help='labeling params',
type=pathlib.Path)
args = parser.parse_args()
rospy.init_node('test_model_from_gazebo')
test_config = json.load(args.test_config.open("r"))
labeling_params = json.load(args.labeling_params.open("r"))
labeling_state_key = labeling_params['state_key']
# read actions from config
actions = [numpify(a) for a in test_config['actions']]
n_actions = len(actions)
time_steps = np.arange(n_actions + 1)
fwd_model, _ = dynamics_utils.load_generic_model(args.fwd_model_dir)
service_provider = GazeboServices()
service_provider.setup_env(
verbose=0,
real_time_rate=0,
max_step_size=fwd_model.data_collection_params['max_step_size'])
environment = fwd_model.scenario.get_environment(
params=fwd_model.data_collection_params)
start_state = fwd_model.scenario.get_state()
start_state = make_dict_tf_float32(start_state)
start_states = [start_state]
expanded_actions = [[actions]]
predicted_states = predict(fwd_model, environment, start_states,
expanded_actions, n_actions, 1, 1)
scenario = fwd_model.scenario
actual_states_lists = execute(service_provider, scenario, start_states,
expanded_actions)
visualize(scenario, environment, actual_states_lists, actions,
predicted_states, labeling_state_key, time_steps)
def test_classifier(classifier_model_dir: pathlib.Path,
fwd_model_dir: List[pathlib.Path],
n_actions: int,
saved_state: Optional[pathlib.Path],
generate_actions: Callable):
fwd_model, _ = dynamics_utils.load_generic_model([pathlib.Path(p) for p in fwd_model_dir])
classifier: NNClassifierWrapper = classifier_utils.load_generic_model([classifier_model_dir])
service_provider = GazeboServices()
service_provider.setup_env(verbose=0,
real_time_rate=0,
max_step_size=fwd_model.data_collection_params['max_step_size'],
play=False)
if saved_state:
service_provider.restore_from_bag(saved_state)
scenario = fwd_model.scenario
scenario.on_before_get_state_or_execute_action()
# NOTE: perhaps it would make sense to have a "fwd_model" have API for get_env, get_state, sample_action, etc
# because the fwd_model knows it's scenario, and importantly it also knows it's data_collection_params
# which is what we're using here to pass to the scenario methods
params = fwd_model.data_collection_params
environment = numpify(scenario.get_environment(params))
start_state = numpify(scenario.get_state())
start_state_tiled = repeat(start_state, n_actions, axis=0, new_axis=True)
start_states_tiled = add_time_dim(start_state_tiled)
actions = generate_actions(environment, start_state_tiled, scenario, params, n_actions)
environment_tiled = repeat(environment, n_actions, axis=0, new_axis=True)
actions_dict = sequence_of_dicts_to_dict_of_tensors(actions)
actions_dict = add_time_dim(actions_dict)
predictions, _ = fwd_model.propagate_differentiable_batched(environment=environment_tiled,
state=start_states_tiled,
actions=actions_dict)
# Run classifier
state_sequence_length = 2
accept_probabilities, _ = classifier.check_constraint_batched_tf(environment=environment_tiled,
predictions=predictions,
actions=actions_dict,
state_sequence_length=state_sequence_length,
batch_size=n_actions)
# animate over the sampled actions
anim = RvizAnimation(scenario=scenario,
n_time_steps=n_actions,
init_funcs=[init_viz_env],
t_funcs=[
lambda s, e, t: init_viz_env(s, e),
viz_transition_for_model_t({}, fwd_model),
ExperimentScenario.plot_accept_probability_t,
],
)
example = {
'accept_probability': tf.squeeze(accept_probabilities, axis=1),
}
example.update(environment)
example.update(predictions)
example.update(actions_dict)
anim.play(example)