def main():
tf.random.set_seed(0)
np.random.seed(0)
colorama.init(autoreset=True)
np.set_printoptions(linewidth=200, precision=3, suppress=True)
parser = argparse.ArgumentParser()
parser.add_argument("dataset_dirs", type=pathlib.Path, nargs="+")
parser.add_argument("checkpoint", type=pathlib.Path)
parser.add_argument("--mode",
type=str,
choices=['train', 'val', 'test'],
default='val')
args = parser.parse_args()
# TODO: REMOVE ME!
args.mode = 'train'
rospy.init_node("test_as_inverse_model")
test_dataset = DynamicsDatasetLoader(args.dataset_dirs)
test_tf_dataset = test_dataset.get_datasets(mode=args.mode)
filter_model = filter_utils.load_filter([args.checkpoint])
latent_dynamics_model, _ = dynamics_utils.load_generic_model(
[args.checkpoint])
test_as_inverse_model(filter_model, latent_dynamics_model, test_dataset,
test_tf_dataset)
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 load_dataset_and_models(args):
comparison_info = json.load(args.comparison.open("r"))
models = {}
for name, model_info in comparison_info.items():
model_dir = paths_from_json(model_info['model_dir'])
model, _ = dynamics_utils.load_generic_model(model_dir)
models[name] = model
dataset = DynamicsDatasetLoader(args.dataset_dirs)
tf_dataset = dataset.get_datasets(mode=args.mode,
shard=args.shard,
take=args.take)
tf_dataset = batch_tf_dataset(tf_dataset, 1)
return tf_dataset, dataset, models
def viz_dataset(
dataset_dirs: List[pathlib.Path],
checkpoint: pathlib.Path,
mode: str,
viz_func: Callable,
use_gt_rope: bool,
**kwargs,
):
test_dataset = DynamicsDatasetLoader(dataset_dirs, use_gt_rope=use_gt_rope)
test_tf_dataset = test_dataset.get_datasets(mode=mode)
test_tf_dataset = batch_tf_dataset(test_tf_dataset, 1, drop_remainder=True)
model, _ = dynamics_utils.load_generic_model([checkpoint])
for i, batch in enumerate(test_tf_dataset):
batch.update(test_dataset.batch_metadata)
outputs, _ = model.from_example(batch, training=False)
viz_func(batch, outputs, test_dataset, model)
def main():
colorama.init(autoreset=True)
parser = argparse.ArgumentParser()
parser.add_argument("results_dir",
type=pathlib.Path,
help='dir containing *_metrics.json.gz')
parser.add_argument("plan_idx", type=int)
args = parser.parse_args()
rospy.init_node("postprocess_result")
with (args.results_dir / 'metadata.json').open('r') as metadata_file:
metadata_str = metadata_file.read()
metadata = json.loads(metadata_str)
planner_params = metadata['planner_params']
fwd_model_dirs = [pathlib.Path(p) for p in planner_params['fwd_model_dir']]
fwd_model, _ = dynamics_utils.load_generic_model(fwd_model_dirs)
scenario = fwd_model.scenario
classifier_model_dir = pathlib.Path(planner_params['classifier_model_dir'])
classifier = classifier_utils.load_generic_model(classifier_model_dir,
scenario=scenario)
metrics_filename = args.results_dir / f"{args.plan_idx}_metrics.json.gz"
with gzip.open(metrics_filename, 'rb') as f:
data_str = f.read()
datum = json.loads(data_str.decode("utf-8"))
steps = datum['steps']
goal = datum['goal']
first_step = steps[0]
environment = numpify(first_step['planning_query']['environment'])
all_output_actions = []
for step in steps:
if step['type'] == 'executed_plan':
actions = postprocess_step(scenario, fwd_model, classifier,
environment, step, goal, planner_params)
all_output_actions.extend(actions)
elif step['type'] == 'executed_recovery':
actions = [step['recovery_action']]
all_output_actions.extend(actions)
# viz the whole thing
start_state = steps[0]['planning_result']['path'][0]
final_states = fwd_model.propagate(environment, start_state, actions)
T = len(final_states)
for t, s_t in enumerate(final_states):
scenario.plot_state_rviz(s_t,
idx=t,
label='smoothed',
color='#00ff0099')
if t < T - 1:
scenario.plot_action_rviz(s_t,
actions[t],
idx=t,
color="#ffffff99",
label='smoothed')
sleep(0.02)
# Save the output actions
outfilename = args.results_dir / f"{args.plan_idx}_smoothed.json"
with outfilename.open("w") as outfile:
json.dump(listify(all_output_actions), outfile)
print(Fore.GREEN + f"Wrote {outfilename}" + Fore.RESET)
def load_fwd_model(planner_params):
fwd_model_dirs = paths_from_json(planner_params['fwd_model_dir'])
fwd_model, _ = dynamics_utils.load_generic_model(fwd_model_dirs)
return fwd_model
def make_recovery_dataset_from_params_dict(dataset_dir: pathlib.Path,
fwd_model_dir,
classifier_model_dir: pathlib.Path,
labeling_params: Dict,
outdir: pathlib.Path,
batch_size: int,
use_gt_rope: bool,
start_at: Optional = None,
stop_at: Optional = None):
# append "best_checkpoint" before loading
classifier_model_dir = classifier_model_dir / 'best_checkpoint'
if not isinstance(fwd_model_dir, List):
fwd_model_dir = [fwd_model_dir]
fwd_model_dir = [p / 'best_checkpoint' for p in fwd_model_dir]
np.random.seed(0)
tf.random.set_seed(0)
dynamics_hparams = hjson.load((dataset_dir / 'hparams.hjson').open('r'))
fwd_model, _ = dynamics_utils.load_generic_model(fwd_model_dir)
dataset = DynamicsDatasetLoader([dataset_dir], use_gt_rope=use_gt_rope)
outdir.mkdir(exist_ok=True)
print(Fore.GREEN + f"Making recovery dataset {outdir.as_posix()}")
new_hparams_filename = outdir / 'hparams.hjson'
recovery_dataset_hparams = dynamics_hparams
scenario = fwd_model.scenario
if not isinstance(classifier_model_dir, List):
classifier_model_dir = [classifier_model_dir]
classifier_model = classifier_utils.load_generic_model(
classifier_model_dir, scenario)
recovery_dataset_hparams['dataset_dir'] = dataset_dir
recovery_dataset_hparams['fwd_model_dir'] = fwd_model_dir
recovery_dataset_hparams['classifier_model'] = classifier_model_dir
recovery_dataset_hparams['fwd_model_hparams'] = fwd_model.hparams
recovery_dataset_hparams['labeling_params'] = labeling_params
recovery_dataset_hparams['state_keys'] = fwd_model.state_keys
recovery_dataset_hparams['action_keys'] = fwd_model.action_keys
recovery_dataset_hparams['start-at'] = start_at
recovery_dataset_hparams['stop-at'] = stop_at
my_hdump(recovery_dataset_hparams,
new_hparams_filename.open("w"),
indent=2)
outdir.mkdir(parents=True, exist_ok=True)
start_at = progress_point(start_at)
stop_at = progress_point(stop_at)
modes = ['train', 'val', 'test']
for mode in modes:
if start_at is not None and modes.index(mode) < modes.index(
start_at[0]):
continue
if stop_at is not None and modes.index(mode) > modes.index(stop_at[0]):
continue
tf_dataset_for_mode = dataset.get_datasets(mode=mode)
full_output_directory = outdir / mode
full_output_directory.mkdir(parents=True, exist_ok=True)
# figure out that record_idx to start at
record_idx = count_up_to_next_record_idx(full_output_directory)
# FIXME: start_at is not implemented correctly in the sense that it shouldn't be the same
# across train/val/test
for out_example in generate_recovery_examples(
tf_dataset=tf_dataset_for_mode,
modes=modes,
mode=mode,
fwd_model=fwd_model,
classifier_model=classifier_model,
dataset=dataset,
labeling_params=labeling_params,
batch_size=batch_size,
start_at=start_at,
stop_at=stop_at):
# FIXME: is there an extra time/batch dimension?
for batch_idx in range(out_example['traj_idx'].shape[0]):
out_example_b = index_dict_of_batched_tensors_tf(
out_example, batch_idx)
# # BEGIN DEBUG
# from link_bot_data.visualization import init_viz_env, recovery_transition_viz_t, init_viz_action
# from copy import deepcopy
#
# viz_out_example_b = deepcopy(out_example_b)
# recovery_probability = compute_recovery_probabilities(viz_out_example_b['accept_probabilities'],
# labeling_params['n_action_samples'])
# viz_out_example_b['recovery_probability'] = recovery_probability
# anim = RvizAnimation(scenario=scenario,
# n_time_steps=labeling_params['action_sequence_horizon'],
# init_funcs=[init_viz_env,
# init_viz_action(dataset.scenario_metadata, fwd_model.action_keys,
# fwd_model.state_keys),
# ],
# t_funcs=[init_viz_env,
# recovery_transition_viz_t(dataset.scenario_metadata,
# fwd_model.state_keys),
# lambda s, e, t: scenario.plot_recovery_probability_t(e, t),
# ])
# anim.play(viz_out_example_b)
# # END DEBUG
tf_write_example(full_output_directory, out_example_b,
record_idx)
record_idx += 1
return outdir
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)
def make_classifier_dataset_from_params_dict(dataset_dir: pathlib.Path,
fwd_model_dir: List[pathlib.Path],
labeling_params: Dict,
outdir: pathlib.Path,
use_gt_rope: bool,
visualize: bool,
take: Optional[int] = None,
batch_size: Optional[int] = None,
start_at: Optional[int] = None,
stop_at: Optional[int] = None):
# append "best_checkpoint" before loading
if not isinstance(fwd_model_dir, List):
fwd_model_dir = [fwd_model_dir]
fwd_model_dir = [p / 'best_checkpoint' for p in fwd_model_dir]
dynamics_hparams = hjson.load((dataset_dir / 'hparams.hjson').open('r'))
fwd_models, _ = dynamics_utils.load_generic_model(fwd_model_dir)
dataset = DynamicsDatasetLoader([dataset_dir], use_gt_rope=use_gt_rope)
new_hparams_filename = outdir / 'hparams.hjson'
classifier_dataset_hparams = dynamics_hparams
classifier_dataset_hparams['dataset_dir'] = dataset_dir.as_posix()
classifier_dataset_hparams['fwd_model_hparams'] = fwd_models.hparams
classifier_dataset_hparams['labeling_params'] = labeling_params
classifier_dataset_hparams['true_state_keys'] = dataset.state_keys
classifier_dataset_hparams['predicted_state_keys'] = fwd_models.state_keys
classifier_dataset_hparams['action_keys'] = dataset.action_keys
classifier_dataset_hparams['scenario_metadata'] = dataset.hparams[
'scenario_metadata']
classifier_dataset_hparams['start-at'] = start_at
classifier_dataset_hparams['stop-at'] = stop_at
my_hdump(classifier_dataset_hparams,
new_hparams_filename.open("w"),
indent=2)
# because we're currently making this dataset, we can't call "get_dataset" but we can still use it to visualize
classifier_dataset_for_viz = ClassifierDatasetLoader(
[outdir], use_gt_rope=use_gt_rope)
t0 = perf_counter()
total_example_idx = 0
for mode in ['train', 'val', 'test']:
tf_dataset = dataset.get_datasets(mode=mode, take=take)
full_output_directory = outdir / mode
full_output_directory.mkdir(parents=True, exist_ok=True)
out_examples_gen = generate_classifier_examples(
fwd_models, tf_dataset, dataset, labeling_params, batch_size)
for out_examples in out_examples_gen:
for out_examples_for_start_t in out_examples:
actual_batch_size = out_examples_for_start_t['traj_idx'].shape[
0]
for batch_idx in range(actual_batch_size):
out_example_b = index_dict_of_batched_tensors_tf(
out_examples_for_start_t, batch_idx)
if out_example_b['time_idx'].ndim == 0:
continue
if visualize:
add_label(out_example_b, labeling_params['threshold'])
classifier_dataset_for_viz.anim_transition_rviz(
out_example_b)
tf_write_example(full_output_directory, out_example_b,
total_example_idx)
rospy.loginfo_throttle(
10,
f"Examples: {total_example_idx:10d}, Time: {perf_counter() - t0:.3f}"
)
total_example_idx += 1
return outdir