def viz_func(batch, predictions, test_dataset: DynamicsDatasetLoader):
""" we assume batch size of 1 """
test_dataset.scenario.plot_environment_rviz(remove_batch(batch))
anim = RvizAnimationController(np.arange(test_dataset.steps_per_traj))
while not anim.done:
t = anim.t()
actual_t = numpify(
remove_batch(
test_dataset.scenario.index_time_batched_predicted(batch, t)))
action_t = numpify(
remove_batch(
test_dataset.scenario.index_time_batched_predicted(batch, t)))
test_dataset.scenario.plot_state_rviz(actual_t,
label='actual',
color='red')
test_dataset.scenario.plot_action_rviz(actual_t,
action_t,
color='gray')
prediction_t = remove_batch(
test_dataset.scenario.index_time_batched_predicted(predictions, t))
test_dataset.scenario.plot_state_rviz(prediction_t,
label='predicted',
color='blue')
anim.step()
def filter_differentiable(self, environment: Dict, state: Optional[Dict], observation: Dict) -> Tuple[Dict, Dict]:
net_inputs = {}
net_inputs.update(environment)
net_inputs.update(observation)
if state is not None:
net_inputs.update(state)
net_inputs = add_batch(net_inputs)
net_inputs = make_dict_tf_float32(net_inputs)
mean_state, stdev_state = self.from_example(net_inputs, training=False)
mean_state = remove_batch(mean_state)
stdev_state = remove_batch(stdev_state)
return mean_state, stdev_state
def test_is_reconverging(self):
batch_is_reconverging_output = is_reconverging(
tf.constant([[1, 0, 0, 1], [1, 1, 1, 0], [1, 0, 0, 0]],
tf.int64)).numpy()
self.assertTrue(batch_is_reconverging_output[0])
self.assertFalse(batch_is_reconverging_output[1])
self.assertFalse(batch_is_reconverging_output[2])
self.assertTrue(
remove_batch(is_reconverging(tf.constant([[1, 0, 0, 1]],
tf.int64))).numpy())
self.assertFalse(
remove_batch(is_reconverging(tf.constant([[1, 0, 0, 0]],
tf.int64))).numpy())
self.assertFalse(
remove_batch(is_reconverging(tf.constant([[1, 0, 1, 0]],
tf.int64))).numpy())
def viz_func(self, batch, outputs, test_dataset: DynamicsDatasetLoader):
""" we assume batch size of 1 """
test_dataset.scenario.plot_environment_rviz(remove_batch(batch))
anim = RvizAnimationController(np.arange(test_dataset.steps_per_traj))
while not anim.done:
t = anim.t()
if self.first_latent_state is None:
self.first_latent_state = outputs['z'][0, 0]
pass
m = Marker()
m.header.frame_id = 'world'
m.header.stamp = rospy.Time.now()
m.type = Marker.SPHERE
m.action = Marker.MODIFY
m.scale.x = 0.01
m.scale.y = 0.01
m.scale.z = 0.01
m.color.r = 0.8
m.color.g = 0.2
m.color.b = 0.8
m.color.a = 0.8
m.id = self.idx
m.ns = 'latent state'
m.pose.position.x = (outputs['z'][0, 0, 0] -
self.first_latent_state[0]) * 10
m.pose.position.y = (outputs['z'][0, 0, 1] -
self.first_latent_state[1]) * 10
m.pose.position.z = (outputs['z'][0, 0, 2] -
self.first_latent_state[2]) * 10
m.pose.orientation.w = 1
self.latent_state_pub.publish(m)
e_t = numpify(
remove_batch(
test_dataset.scenario.index_time_batched_predicted(
batch, t)))
test_dataset.scenario.plot_state_rviz(e_t,
label='actual',
color='red')
test_dataset.scenario.plot_action_rviz(e_t, e_t, color='gray')
self.idx += 1
anim.step()
def propagate_differentiable(self, environment: Dict, start_state: Dict, actions: List[Dict]) -> Tuple[
List[Dict], List[Dict]]:
# add time dimension of size 1
net_inputs = {k: tf.expand_dims(start_state[k], axis=0) for k in self.state_keys}
net_inputs.update(environment)
net_inputs.update(sequence_of_dicts_to_dict_of_tensors(actions))
net_inputs = add_batch(net_inputs)
net_inputs = make_dict_tf_float32(net_inputs)
# the network returns a dictionary where each value is [T, n_state]
# which is what you'd want for training, but for planning and execution and everything else
# it is easier to deal with a list of states where each state is a dictionary
mean_predictions, stdev_predictions = self.from_example(net_inputs, training=False)
mean_predictions = remove_batch(mean_predictions)
stdev_predictions = remove_batch(stdev_predictions)
mean_predictions = dict_of_sequences_to_sequence_of_dicts_tf(mean_predictions)
stdev_predictions = dict_of_sequences_to_sequence_of_dicts_tf(stdev_predictions)
return mean_predictions, stdev_predictions
def viz_main(args):
dataset_dirs = args.dataset_dirs
checkpoint = args.checkpoint
trial_path, params = load_trial(checkpoint.parent.absolute())
dataset = DynamicsDatasetLoader(dataset_dirs)
scenario = dataset.scenario
tf_dataset = dataset.get_datasets(mode='val')
tf_dataset = batch_tf_dataset(tf_dataset,
batch_size=1,
drop_remainder=True)
model = CFM(hparams=params, batch_size=1, scenario=scenario)
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, args.checkpoint, max_to_keep=1)
status = ckpt.restore(manager.latest_checkpoint).expect_partial()
if manager.latest_checkpoint:
print(Fore.CYAN + "Restored from {}".format(manager.latest_checkpoint))
status.assert_existing_objects_matched()
else:
raise RuntimeError("Failed to restore!!!")
for example_idx, example in enumerate(tf_dataset):
stepper = RvizAnimationController(n_time_steps=dataset.steps_per_traj)
for t in range(dataset.steps_per_traj):
output = model(
model.preprocess_no_gradient(example, training=False))
actual_t = numpify(
remove_batch(scenario.index_time_batched_predicted(example,
t)))
action_t = numpify(
remove_batch(scenario.index_time_batched_predicted(example,
t)))
scenario.plot_state_rviz(actual_t, label='actual', color='red')
scenario.plot_action_rviz(actual_t, action_t, color='gray')
prediction_t = remove_batch(
scenario.index_time_batched_predicted(output, t))
scenario.plot_state_rviz(prediction_t,
label='predicted',
color='blue')
stepper.step()
def sample(self, environment: Dict, state: Dict):
input_dict = environment
input_dict.update({add_predicted(k): tf.expand_dims(v, axis=0) for k, v in state.items()})
input_dict = add_batch(input_dict)
input_dict = {k: tf.cast(v, tf.float32) for k, v in input_dict.items()}
output = self.net.sample(input_dict)
output = remove_batch(output)
output = numpify(output)
return output
def check_constraint_tf(self,
environment: Dict,
states_sequence: List[Dict],
actions: List[Dict]):
environment = add_batch(environment)
states_sequence_dict = sequence_of_dicts_to_dict_of_tensors(states_sequence)
states_sequence_dict = add_batch(states_sequence_dict)
state_sequence_length = len(states_sequence)
actions_dict = sequence_of_dicts_to_dict_of_tensors(actions)
actions_dict = add_batch(actions_dict)
mean_probabilities, stdev_probabilities = self.check_constraint_batched_tf(environment=environment,
predictions=states_sequence_dict,
actions=actions_dict,
batch_size=1,
state_sequence_length=state_sequence_length)
mean_probabilities = remove_batch(mean_probabilities)
stdev_probabilities = remove_batch(stdev_probabilities)
return mean_probabilities, stdev_probabilities
def viz_example(batch, outputs, test_dataset: DynamicsDatasetLoader, model):
test_dataset.scenario.plot_environment_rviz(remove_batch(batch))
anim = RvizAnimationController(np.arange(test_dataset.steps_per_traj))
while not anim.done:
t = anim.t()
actual_t = test_dataset.index_time_batched(batch, t)
test_dataset.scenario.plot_state_rviz(actual_t,
label='actual',
color='red')
test_dataset.scenario.plot_action_rviz(actual_t,
actual_t,
color='gray')
prediction_t = test_dataset.index_time_batched(outputs, t)
test_dataset.scenario.plot_state_rviz(prediction_t,
label='predicted',
color='blue')
anim.step()
def propagate_differentiable(self, full_env: np.ndarray,
full_env_origin: np.ndarray, res: float,
start_state: Dict[str, np.ndarray],
actions: tf.Variable) -> List[Dict]:
"""
:param full_env: (H, W)
:param full_env_origin: (2)
:param res: scalar
:param start_state: each value in the dictionary should be of shape (batch, n_state)
:param actions: (T, 2)
:return: states: each value in the dictionary should be a of shape [batch, T+1, n_state)
"""
test_x = {
# shape: T, 2
'action':
tf.convert_to_tensor(actions, dtype=tf.float32),
# shape: 1
'res':
tf.convert_to_tensor(res, dtype=tf.float32),
# shape: H, W
'full_env/env':
tf.convert_to_tensor(full_env, dtype=tf.float32),
# shape: 2
'full_env/origin':
tf.convert_to_tensor(full_env_origin, dtype=tf.float32),
# scalar
'full_env/res':
tf.convert_to_tensor(res, dtype=tf.float32),
}
for state_key, v in start_state.items():
# handles conversion from double -> float
start_state = tf.convert_to_tensor(v, dtype=tf.float32)
start_state_with_time_dim = tf.expand_dims(start_state, axis=0)
test_x[state_key] = start_state_with_time_dim
test_x = add_batch(test_x)
predictions = self.net((test_x, False))
predictions = remove_batch(predictions)
predictions = dict_of_sequences_to_sequence_of_dicts_tf(predictions)
return predictions
def get_rgbd(self):
color_msg: Image = self.color_image_listener.get()
depth_msg = self.depth_image_listener.get()
depth = np.expand_dims(ros_numpy.numpify(depth_msg), axis=-1)
bgr = ros_numpy.numpify(color_msg)
rgb = np.flip(bgr, axis=2)
# NaN Depths means out of range, so clip to the max range
depth = np.clip(np.nan_to_num(depth, nan=KINECT_MAX_DEPTH), 0,
KINECT_MAX_DEPTH)
rgbd = np.concatenate([rgb, depth], axis=2)
box = tf.convert_to_tensor([
self.crop_region['min_y'] / rgb.shape[0],
self.crop_region['min_x'] / rgb.shape[1],
self.crop_region['max_y'] / rgb.shape[0],
self.crop_region['max_x'] / rgb.shape[1]
],
dtype=tf.float32)
# this operates on a batch
rgbd_cropped = tf.image.crop_and_resize(
image=tf.expand_dims(rgbd, axis=0),
boxes=tf.expand_dims(box, axis=0),
box_indices=[0],
crop_size=[self.IMAGE_H, self.IMAGE_W])
rgbd_cropped = remove_batch(rgbd_cropped)
def _debug_show_image(_rgb_depth_cropped):
import matplotlib.pyplot as plt
plt.imshow(tf.cast(_rgb_depth_cropped[:, :, :3], tf.int32))
plt.show()
# BEGIN DEBUG
# _debug_show_image(rgbd_cropped)
# END DEBUG
return rgbd_cropped.numpy()
def index_time_batched(self, example_batched, t: int):
e_t = numpify(
remove_batch(
index_time_batched(example_batched, self.time_indexed_keys,
t)))
return e_t
def test_as_inverse_model(filter_model, latent_dynamics_model, test_dataset,
test_tf_dataset):
scenario = test_dataset.scenario
shooting_method = ShootingMethod(fwd_model=latent_dynamics_model,
classifier_model=None,
scenario=scenario,
params={'n_samples': 1000})
trajopt = TrajectoryOptimizer(fwd_model=latent_dynamics_model,
classifier_model=None,
scenario=scenario,
params={
"iters": 100,
"length_alpha": 0,
"goal_alpha": 1000,
"constraints_alpha": 0,
"action_alpha": 0,
"initial_learning_rate": 0.0001,
})
s_color_viz_pub = rospy.Publisher("s_state_color_viz",
Image,
queue_size=10,
latch=True)
s_next_color_viz_pub = rospy.Publisher("s_next_state_color_viz",
Image,
queue_size=10,
latch=True)
image_diff_viz_pub = rospy.Publisher("image_diff_viz",
Image,
queue_size=10,
latch=True)
action_horizon = 1
initial_actions = []
total_errors = []
for example_idx, example in enumerate(test_tf_dataset):
stepper = RvizAnimationController(
n_time_steps=test_dataset.steps_per_traj)
for t in range(test_dataset.steps_per_traj - 1):
print(example_idx)
environment = {}
current_observation = remove_batch(
scenario.index_observation_time_batched(add_batch(example), t))
for j in range(action_horizon):
left_gripper_position = [0, 0, 0]
right_gripper_position = [0, 0, 0]
initial_action = {
'left_gripper_position': left_gripper_position,
'right_gripper_position': right_gripper_position,
}
initial_actions.append(initial_action)
goal_observation = {
k: example[k][1]
for k in filter_model.obs_keys
}
planning_query = PlanningQuery(start=current_observation,
goal=goal_observation,
environment=environment,
seed=1)
planning_result = shooting_method.plan(planning_query)
actions = planning_result.actions
planned_path = planning_result.latent_path
true_action = numpify(
{k: example[k][0]
for k in latent_dynamics_model.action_keys})
for j in range(action_horizon):
optimized_action = actions[j]
# optimized_action = {
# 'left_gripper_position': current_observation['left_gripper'],
# 'right_gripper_position': current_observation['right_gripper'],
# }
true_action = numpify({
k: example[k][j]
for k in latent_dynamics_model.action_keys
})
# Visualize
s = numpify(
remove_batch(
scenario.index_observation_time_batched(
add_batch(example), 0)))
s.update(
numpify(
remove_batch(
scenario.index_observation_features_time_batched(
add_batch(example), 0))))
s_next = numpify(
remove_batch(
scenario.index_observation_time_batched(
add_batch(example), 1)))
s_next.update(
numpify(
remove_batch(
scenario.index_observation_features_time_batched(
add_batch(example), 1))))
scenario.plot_state_rviz(s, label='t', color="#ff000055", id=1)
scenario.plot_state_rviz(s_next,
label='t+1',
color="#aa222255",
id=2)
# scenario.plot_action_rviz(s, optimized_action, label='inferred', color='#00ff00', id=1)
# scenario.plot_action_rviz(s, true_action, label='true', color='#ee770055', id=2)
publish_color_image(s_color_viz_pub, s['rgbd'][:, :, :3])
publish_color_image(s_next_color_viz_pub,
s_next['rgbd'][:, :, :3])
diff = s['rgbd'][:, :, :3] - s_next['rgbd'][:, :, :3]
publish_color_image(image_diff_viz_pub, diff)
# Metrics
total_error = 0
for v1, v2 in zip(optimized_action.values(),
true_action.values()):
total_error += -np.dot(v1, v2)
total_errors.append(total_error)
stepper.step()
if example_idx > 100:
break
print(np.min(total_errors))
print(np.max(total_errors))
print(np.mean(total_errors))
plt.xlabel("total error (meter-ish)")
plt.hist(total_errors, bins=np.linspace(0, 2, 20))
plt.show()
def viz(data_filename, fps, no_plot, save):
rospy.init_node("compare_models")
# Load the results
base_folder = data_filename.parent
with gzip.open(data_filename, "rb") as data_file:
data_str = data_file.read()
saved_data = json.loads(data_str.decode("utf-8"))
all_metrics = {}
for example_idx, datum in enumerate(saved_data):
print(example_idx)
# use the first (or any) model data to get the ground truth and
dataset_element = numpify(datum.pop("dataset_element"))
environment = numpify(datum.pop("environment"))
action_keys = datum.pop("action_keys")
actions = {k: dataset_element[k] for k in action_keys}
models_viz_info = {}
n_models = len(datum)
time_steps = np.arange(datum.pop('time_steps'))
for model_name, data_for_model in datum.items():
scenario = get_scenario(data_for_model['scenario'])
# Metrics
metrics_for_model = {}
predictions = numpify(data_for_model['predictions'])
predictions.pop('stdev')
metrics = scenario.dynamics_metrics_function(dataset_element, predictions)
loss = scenario.dynamics_loss_function(dataset_element, predictions)
metrics['loss'] = loss
for metric_name, metric_value in metrics.items():
if metric_name not in metrics_for_model:
metrics_for_model[metric_name] = []
metrics_for_model[metric_name].append(metric_value.numpy())
for metric_name, metric_values in metrics_for_model.items():
mean_metric_value = float(np.mean(metric_values))
if model_name not in all_metrics:
all_metrics[model_name] = {}
if metric_name not in all_metrics[model_name]:
all_metrics[model_name][metric_name] = []
all_metrics[model_name][metric_name].append(mean_metric_value)
models_viz_info[model_name] = (scenario, predictions)
if not no_plot and not save:
# just use whatever the latest scenario was, it shouldn't matter which we use
scenario.plot_environment_rviz(remove_batch(environment))
anim = RvizAnimationController(time_steps)
while not anim.done:
t = anim.t()
actual_t = remove_batch(scenario.index_state_time(dataset_element, t))
action_t = remove_batch(scenario.index_action_time(actions, t))
scenario.plot_state_rviz(actual_t, label='actual', color='#0000ff88')
scenario.plot_action_rviz(actual_t, action_t, color='gray')
for model_idx, (model_name, viz_info) in enumerate(models_viz_info.items()):
scenario_i, predictions = viz_info
prediction_t = remove_batch(scenario_i.index_state_time(predictions, t))
color = cm.jet(model_idx / n_models)
scenario_i.plot_state_rviz(prediction_t, label=model_name, color=color)
anim.step()
metrics_by_model = {}
for model_name, metrics_for_model in all_metrics.items():
for metric_name, metric_values in metrics_for_model.items():
if metric_name not in metrics_by_model:
metrics_by_model[metric_name] = {}
metrics_by_model[metric_name][model_name] = metric_values
with (base_folder / 'metrics_tables.txt').open("w") as metrics_file:
for metric_name, metric_by_model in metrics_by_model.items():
headers = ["Model", "min", "max", "mean", "median", "std"]
table_data = []
for model_name, metric_values in metric_by_model.items():
table_data.append([model_name] + row_stats(metric_values))
print('-' * 90)
print(Style.BRIGHT + metric_name + Style.NORMAL)
table = tabulate(table_data,
headers=headers,
tablefmt='fancy_grid',
floatfmt='6.4f',
numalign='center',
stralign='left')
metrics_file.write(table)
print(table)
print()
print(Style.BRIGHT + f"p-value matrix [{metric_name}]" + Style.NORMAL)
print(dict_to_pvalue_table(metric_by_model))
def filter_no_reconverging(example):
is_close = example['is_close']
return tf.logical_not(remove_batch(is_reconverging(add_batch(is_close))))
def filter_only_reconverging(example):
is_close = example['is_close']
return remove_batch(is_reconverging(add_batch(is_close)))
def index_time(e: Dict, time_indexed_keys: List[str], t: int):
return remove_batch(index_time_batched(add_batch(e), time_indexed_keys, t))
def visualize_dataset(args, classifier_dataset):
tf_dataset = classifier_dataset.get_datasets(mode=args.mode, take=args.take)
tf_dataset = tf_dataset.batch(1)
iterator = iter(tf_dataset)
t0 = perf_counter()
reconverging_count = 0
positive_count = 0
negative_count = 0
count = 0
stdevs = []
labels = []
stdevs_for_negative = []
stdevs_for_positive = []
for i, example in enumerate(progressbar(tf_dataset, widgets=base_dataset.widgets)):
example = remove_batch(example)
is_close = example['is_close'].numpy().squeeze()
count += is_close.shape[0]
n_close = np.count_nonzero(is_close[-1])
n_far = is_close.shape[0] - n_close
positive_count += n_close
negative_count += n_far
reconverging = n_far > 0 and is_close[-1]
if args.only_reconverging and not reconverging:
continue
if args.only_negative and np.any(is_close[1:]):
continue
if args.only_positive and not np.any(is_close[1:]):
continue
# print(f"Example {i}, Trajectory #{int(example['traj_idx'])}")
if count == 0:
print_dict(example)
if reconverging:
reconverging_count += 1
# Print statistics intermittently
if count % 1000 == 0:
print_stats_and_timing(args, count, reconverging_count, negative_count, positive_count)
#############################
# Show Visualization
#############################
if args.display_type == 'just_count':
continue
elif args.display_type == '3d':
# print(example['is_close'])
if example['is_close'][0] == 0:
continue
classifier_dataset.anim_transition_rviz(example)
elif args.display_type == 'stdev':
for t in range(1, classifier_dataset.horizon):
stdev_t = example[add_predicted('stdev')][t, 0].numpy()
label_t = example['is_close'][t]
stdevs.append(stdev_t)
labels.append(label_t)
if label_t > 0.5:
stdevs_for_positive.append(stdev_t)
else:
stdevs_for_negative.append(stdev_t)
else:
raise NotImplementedError()
total_dt = perf_counter() - t0
if args.display_type == 'stdev':
print(f"p={stats.f_oneway(stdevs_for_negative, stdevs_for_positive)[1]}")
plt.figure()
plt.title(" ".join([str(d.name) for d in args.dataset_dirs]))
bins = plt.hist(stdevs_for_negative, label='negative examples', alpha=0.8, density=True)[1]
plt.hist(stdevs_for_positive, label='positive examples', alpha=0.8, bins=bins, density=True)
plt.ylabel("count")
plt.xlabel("stdev")
plt.legend()
plt.show()
print_stats_and_timing(args, count, reconverging_count, negative_count, positive_count, total_dt)
def _gen(e):
example = next(goal_dataset_iterator)
example_t = dataset.index_time_batched(example_batched=add_batch(example), t=1)
goal = remove_batch(example_t)
return goal