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 debug_rviz(self, input_dict: Dict, debug_info_seq: List[Tuple]):
import pickle
from merrrt_visualization.rviz_animation_controller import RvizSimpleStepper
from link_bot_pycommon.bbox_visualization import grid_to_bbox
from moonshine.moonshine_utils import numpify
from moonshine.indexing import index_dict_of_batched_tensors_tf, index_time_with_metadata
from link_bot_pycommon.grid_utils import environment_to_occupancy_msg, send_occupancy_tf
import numpy as np
from time import time
debug_filename = f'debug_{int(time())}.pkl'
print(f"Saving debug info to {debug_filename}")
with open(debug_filename, "wb") as debug_file:
pickle.dump({'input_dict': input_dict, 'debug_info': debug_info_seq}, debug_file)
stepper = RvizSimpleStepper()
batch_size = input_dict.pop("batch_size").numpy().astype(np.int32)
input_dict.pop("time")
for b in range(batch_size):
example = index_dict_of_batched_tensors_tf(input_dict, b)
for t, debug_info_t in enumerate(debug_info_seq):
state_t, local_env_origin_t, local_env_t, local_voxel_grid_t = debug_info_t
for i, state_component_k_voxel_grid in enumerate(tf.transpose(local_voxel_grid_t, [4, 0, 1, 2, 3])):
raster_dict = {
'env': tf.clip_by_value(state_component_k_voxel_grid[b], 0, 1),
'origin': local_env_origin_t[b].numpy(),
'res': input_dict['res'][b].numpy(),
}
raster_msg = environment_to_occupancy_msg(raster_dict, frame='local_occupancy')
self.raster_debug_pubs[i].publish(raster_msg)
local_env_dict = {
'env': local_env_t[b],
'origin': local_env_origin_t[b].numpy(),
'res': input_dict['res'][b].numpy(),
}
send_occupancy_tf(self.scenario.tf.tf_broadcaster, local_env_dict, frame='local_occupancy')
pred_t = index_time_with_metadata({}, example, self.pred_state_keys, t)
action_t = {k: example[k][t] for k in self.action_keys}
self.scenario.plot_state_rviz(numpify(pred_t), label='predicted', color='#0000ffff')
if action_t is not None:
self.scenario.plot_action_rviz(numpify(pred_t), numpify(action_t), label='action',
color='#0000ffff')
# # Ground-Truth
# true_t = index_time_with_metadata({}, example, self.true_state_keys, t)
# self.scenario.plot_state_rviz(numpify(true_t), label='actual', color='#ff0000ff', scale=1.1)
# label_t = example['is_close'][1]
# self.scenario.plot_is_close(label_t)
bbox_msg = grid_to_bbox(rows=self.local_env_h_rows,
cols=self.local_env_w_cols,
channels=self.local_env_c_channels,
resolution=example['res'].numpy())
bbox_msg.header.frame_id = 'local_occupancy'
self.local_env_bbox_pub.publish(bbox_msg)
stepper.step()
def plan_internal(self, planning_query: PlanningQuery) -> PlanningResult:
start_planning_time = perf_counter()
action_rng = np.random.RandomState(planning_query.seed)
random_actions = self.scenario.sample_action_sequences(environment=planning_query.environment,
state=planning_query.start,
action_params=self.action_params,
n_action_sequences=self.n_samples,
action_sequence_length=1,
validate=True,
action_rng=action_rng)
random_actions = [sequence_of_dicts_to_dict_of_tensors(a) for a in random_actions]
random_actions = sequence_of_dicts_to_dict_of_tensors(random_actions)
environment_batched = {k: tf.stack([v] * self.n_samples, axis=0) for k, v in planning_query.environment.items()}
start_state_batched = {k: tf.expand_dims(tf.stack([v] * self.n_samples, axis=0), axis=1) for k, v in
planning_query.start.items()}
mean_predictions, _ = self.fwd_model.propagate_differentiable_batched(environment=environment_batched,
state=start_state_batched,
actions=random_actions)
final_states = {k: v[:, -1] for k, v in mean_predictions.items()}
goal_state_batched = {k: tf.stack([v] * self.n_samples, axis=0) for k, v in planning_query.goal.items()}
costs = self.scenario.trajopt_distance_to_goal_differentiable(final_states, goal_state_batched)
costs = tf.squeeze(costs)
min_idx = tf.math.argmin(costs, axis=0)
# print(costs)
# print('min idx', min_idx.numpy())
best_indices = tf.argsort(costs)
cmap = cm.Blues
n_to_show = 5
min_cost = costs[best_indices[0]]
max_cost = costs[best_indices[n_to_show]]
for j, i in enumerate(best_indices[:n_to_show]):
s = numpify({k: v[i] for k, v in start_state_batched.items()})
a = numpify({k: v[i][0] for k, v in random_actions.items()})
c = (costs[i] - min_cost) / (max_cost - min_cost)
self.scenario.plot_action_rviz(s, a, label='samples', color=cmap(c), idx1=2 * j, idx2=2 * j + 1)
best_actions = {k: v[min_idx] for k, v in random_actions.items()}
best_prediction = {k: v[min_idx] for k, v in mean_predictions.items()}
best_actions = numpify(dict_of_sequences_to_sequence_of_dicts(best_actions))
best_predictions = numpify(dict_of_sequences_to_sequence_of_dicts(best_prediction))
planning_time = perf_counter() - start_planning_time
planner_status = MyPlannerStatus.Solved
return PlanningResult(status=planner_status, path=best_predictions, actions=best_actions, time=planning_time, tree={})
def plot_3d(args, dataset: DynamicsDatasetLoader, tf_dataset: tf.data.Dataset):
scenario = dataset.scenario
image_diff_viz_pub = rospy.Publisher("image_diff_viz",
Image,
queue_size=10,
latch=True)
for i, example in enumerate(tf_dataset):
if args.start_at is not None and i < args.start_at:
continue
example = numpify(example)
print(f"Example {i}, Trajectory #{int(example['traj_idx'])}")
time_steps = example['time_idx']
anim = RvizAnimationController(time_steps)
while not anim.done:
t = anim.t()
scenario.plot_environment_rviz(example)
example_t = index_time_np(example, dataset.time_indexed_keys, t)
scenario.plot_state_rviz(example_t, label='')
scenario.plot_action_rviz_internal(example_t, label='')
if t < dataset.steps_per_traj - 1:
s_next = index_time_np(example, dataset.time_indexed_keys,
t + 1)
# diff = s['rgbd'][:, :, :3] - s_next['rgbd'][:, :, :3]
# publish_color_image(image_diff_viz_pub, diff)
# this will return when either the animation is "playing" or because the user stepped forward
anim.step()
def on_trial_complete(self, trial_data: Dict, trial_idx: int):
extra_trial_data = {
"planner_params": self.planner_params,
"scenario": self.planner.scenario.simple_name(),
'current_time': int(time()),
'uuid': uuid.uuid4(),
}
trial_data.update(extra_trial_data)
data_filename = self.outdir / f'{trial_idx}_metrics.json.gz'
dummy_proof_write(trial_data, data_filename)
if self.record:
# TODO: maybe make this happen async?
sleep(1)
self.service_provider.stop_record_trial()
self.bag.close()
# compute the current running average success
goal = trial_data['planning_queries'][0].goal
final_actual_state = numpify(trial_data['end_state'])
final_execution_to_goal_error = self.planner.scenario.distance_to_goal(final_actual_state, goal)
self.final_execution_to_goal_errors.append(final_execution_to_goal_error)
goal_threshold = self.planner_params['goal_params']['threshold']
n = len(self.final_execution_to_goal_errors)
success_percentage = np.count_nonzero(np.array(self.final_execution_to_goal_errors) < goal_threshold) / n * 100
update_msg = f"[{self.outdir.name}] Current average success rate {success_percentage:.2f}%"
rospy.loginfo(Fore.LIGHTBLUE_EX + update_msg)
jobkey = self.jobkey(trial_idx)
self.job_chunker.store_result(jobkey, {'data_filename': data_filename})
def plot_recovery(args, scenario, step, metadata):
actual_path = step['execution_result']['path']
action = step['recovery_action']
environment = numpify(step['planning_query']['environment'])
scenario.plot_environment_rviz(environment)
scenario.plot_state_rviz(actual_path[0], idx=1, label='recovery')
scenario.plot_action_rviz(actual_path[0], action, label='recovery')
scenario.plot_state_rviz(actual_path[1], idx=2, label='recovery')
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 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 play(self, example: Any):
example = numpify(example)
for init_func in self.init_funcs:
init_func(self.scenario, example)
controller = RvizAnimationController(n_time_steps=self.n_time_steps)
while not controller.done:
t = controller.t()
for t_func in self.t_funcs:
t_func(self.scenario, example, t)
controller.step()
def evaluate(self, batch_size: int, trial_path: pathlib.Path, take: int):
print(Fore.GREEN + f"Evaluating {trial_path.as_posix()}")
train_metrics = train_test_classifier.eval_main(
dataset_dirs=[self.classifier_dataset_dir],
checkpoint=trial_path / 'best_checkpoint',
mode='train',
threshold=self.threshold,
trials_directory=self.trials_directory,
batch_size=batch_size,
use_gt_rope=False,
take=take,
)
val_metrics = train_test_classifier.eval_main(
dataset_dirs=[self.classifier_dataset_dir],
checkpoint=trial_path / 'best_checkpoint',
mode='val',
threshold=self.threshold,
trials_directory=self.trials_directory,
batch_size=batch_size,
use_gt_rope=False,
take=take)
train_metrics = dict_tools.dict_round(numpify(train_metrics))
val_metrics = dict_tools.dict_round(numpify(val_metrics))
return train_metrics, val_metrics
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 plot_plan(args, scenario, step, metadata, fallback_labeing_params: Dict):
planner_params = metadata['planner_params']
labeling_params = labeling_params_from_planner_params(
planner_params, fallback_labeing_params)
goal = step['planning_query']['goal']
environment = numpify(step['planning_query']['environment'])
planned_path = step['planning_result']['path']
actual_path = step['execution_result']['path']
planned_actions = step['planning_result']['actions']
scenario.reset_planning_viz()
if args.show_tree:
def _draw_tree_function(scenario, tree_json):
print(f"n vertices {len(tree_json['vertices'])}")
for vertex in tree_json['vertices']:
scenario.plot_tree_state(vertex, color='#77777722')
sleep(0.001)
tree_thread = threading.Thread(target=_draw_tree_function,
args=(
scenario,
step['tree_json'],
))
tree_thread.start()
goal_threshold = get_goal_threshold(planner_params)
scenario.animate_evaluation_results(environment=environment,
actual_states=actual_path,
predicted_states=planned_path,
actions=planned_actions,
goal=goal,
goal_threshold=goal_threshold,
labeling_params=labeling_params,
accept_probabilities=None,
horizon=metadata['horizon'])
def reset_rope(self, action_params: Dict):
reset = SetRopeStateRequest()
# TODO: rename this to rope endpoints reset positions or something
reset.left_gripper.x = numpify(
action_params['left_gripper_reset_position'][0])
reset.left_gripper.y = numpify(
action_params['left_gripper_reset_position'][1])
reset.left_gripper.z = numpify(
action_params['left_gripper_reset_position'][2])
reset.right_gripper.x = numpify(
action_params['right_gripper_reset_position'][0])
reset.right_gripper.y = numpify(
action_params['right_gripper_reset_position'][1])
reset.right_gripper.z = numpify(
action_params['right_gripper_reset_position'][2])
self.set_rope_state_srv(reset)
def plot_steps(args, scenario, datum, metadata, fallback_labeing_params: Dict):
planner_params = metadata['planner_params']
goal_threshold = get_goal_threshold(planner_params)
labeling_params = labeling_params_from_planner_params(
planner_params, fallback_labeing_params)
steps = datum['steps']
if len(steps) == 0:
q = datum['planning_queries'][0]
start = q['start']
goal = q['goal']
environment = q['environment']
anim = RvizAnimationController(n_time_steps=1)
scenario.plot_state_rviz(start, label='actual', color='#ff0000aa')
scenario.plot_goal_rviz(goal, goal_threshold)
scenario.plot_environment_rviz(environment)
anim.step()
return
goal = datum['goal']
first_step = steps[0]
environment = numpify(first_step['planning_query']['environment'])
all_actual_states = []
types = []
all_predicted_states = []
all_actions = []
for step_idx, step in enumerate(steps):
if step['type'] == 'executed_plan':
actions = step['planning_result']['actions']
actual_states = step['execution_result']['path']
predicted_states = step['planning_result']['path']
elif step['type'] == 'executed_recovery':
actions = [step['recovery_action']]
actual_states = step['execution_result']['path']
predicted_states = [None, None]
else:
raise NotImplementedError(f"invalid step type {step['type']}")
actions = numpify(actions)
actual_states = numpify(actual_states)
predicted_states = numpify(predicted_states)
all_actions.extend(actions)
types.extend([step['type']] * len(actions))
all_actual_states.extend(actual_states[:-1])
all_predicted_states.extend(predicted_states[:-1])
if args.show_tree and step['type'] == 'executed_plan':
def _draw_tree_function(scenario, tree_json):
print(f"n vertices {len(tree_json['vertices'])}")
for vertex in tree_json['vertices']:
scenario.plot_tree_state(vertex, color='#77777722')
sleep(0.001)
tree_thread = threading.Thread(target=_draw_tree_function,
args=(
scenario,
step['planning_result']['tree'],
))
tree_thread.start()
# but do add the actual final states
all_actual_states.append(actual_states[-1])
all_predicted_states.append(predicted_states[-1])
anim = RvizAnimationController(n_time_steps=len(all_actual_states))
def _type_action_color(type_t: str):
if type_t == 'executed_plan':
return 'b'
elif type_t == 'executed_recovery':
return '#ff00ff'
scenario.reset_planning_viz()
dist_to_goal = np.inf
while not anim.done:
scenario.plot_environment_rviz(environment)
t = anim.t()
s_t = all_actual_states[t]
s_t_pred = all_predicted_states[t]
scenario.plot_state_rviz(s_t, label='actual', color='#ff0000aa')
c = '#0000ffaa'
if len(all_actions) > 0:
if t < anim.max_t:
type_t = types[t]
action_color = _type_action_color(type_t)
scenario.plot_action_rviz(s_t,
all_actions[t],
color=action_color)
else:
type_t = types[t - 1]
action_color = _type_action_color(type_t)
scenario.plot_action_rviz(all_actual_states[t - 1],
all_actions[t - 1],
color=action_color)
if s_t_pred is not None:
scenario.plot_state_rviz(s_t_pred, label='predicted', color=c)
is_close = scenario.compute_label(s_t, s_t_pred, labeling_params)
scenario.plot_is_close(is_close)
else:
scenario.plot_is_close(None)
dist_to_goal = scenario.distance_to_goal(s_t, goal)
actually_at_goal = dist_to_goal < goal_threshold
scenario.plot_goal_rviz(goal, goal_threshold, actually_at_goal)
anim.step()
print(f"final dist to goal {dist_to_goal:.3f}")
def __init__(self,
planner: MyPlanner,
trials: List[int],
verbose: int,
planner_params: Dict,
service_provider: BaseServices,
no_execution: bool,
test_scenes_dir: Optional[pathlib.Path] = None,
save_test_scenes_dir: Optional[pathlib.Path] = None,
):
self.planner = planner
self.scenario = self.planner.scenario
self.scenario.on_before_get_state_or_execute_action()
self.trials = trials
self.planner_params = planner_params
self.verbose = verbose
self.service_provider = service_provider
self.no_execution = no_execution
self.env_rng = np.random.RandomState(0)
self.goal_rng = np.random.RandomState(0)
self.recovery_rng = np.random.RandomState(0)
self.test_scenes_dir = test_scenes_dir
self.save_test_scenes_dir = save_test_scenes_dir
if self.planner_params['recovery']['use_recovery']:
recovery_model_dir = pathlib.Path(self.planner_params['recovery']['recovery_model_dir'])
self.recovery_policy = recovery_policy_utils.load_generic_model(model_dir=recovery_model_dir,
scenario=self.scenario,
rng=self.recovery_rng)
else:
self.recovery_policy = None
self.n_failures = 0
# for saving snapshots of the world
self.link_states_listener = Listener("gazebo/link_states", LinkStates)
# Debugging
if self.verbose >= 2:
self.goal_bbox_pub = rospy.Publisher('goal_bbox', BoundingBox, queue_size=10, latch=True)
bbox_msg = extent_to_bbox(planner_params['goal_params']['extent'])
bbox_msg.header.frame_id = 'world'
self.goal_bbox_pub.publish(bbox_msg)
goal_params = self.planner_params['goal_params']
if goal_params['type'] == 'fixed':
self.goal_generator = lambda e: numpify(goal_params['goal_fixed'])
elif goal_params['type'] == 'random':
self.goal_generator = lambda e: self.scenario.sample_goal(environment=e,
rng=self.goal_rng,
planner_params=self.planner_params)
elif goal_params['type'] == 'dataset':
dataset = DynamicsDatasetLoader([pathlib.Path(goal_params['goals_dataset'])])
tf_dataset = dataset.get_datasets(mode='val')
goal_dataset_iterator = iter(tf_dataset)
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
self.goal_generator = _gen
else:
raise NotImplementedError(f"invalid goal param type {goal_params['type']}")
def propagate(self, environment: Dict, start_states: Dict, actions: List[Dict]) -> Tuple[List[Dict], List[Dict]]:
mean_predictions, stdev_predictions = self.propagate_differentiable(environment, start_states, actions)
return numpify(mean_predictions), numpify(stdev_predictions)
def index_time_np(e: Dict, time_indexed_keys: List[str], t: int):
return numpify(index_time(e, time_indexed_keys, 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 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 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 viz_main(dataset_dirs: List[pathlib.Path],
checkpoint: pathlib.Path,
mode: str,
batch_size: int,
only_errors: bool,
use_gt_rope: bool,
old_compat: bool = False,
**kwargs):
stdev_pub_ = rospy.Publisher("stdev", Float32, queue_size=10)
traj_idx_pub_ = rospy.Publisher("traj_idx_viz", Float32, queue_size=10)
###############
# Model
###############
trials_directory = pathlib.Path('trials').absolute()
trial_path = checkpoint.parent.absolute()
_, params = filepath_tools.create_or_load_trial(
trial_path=trial_path, trials_directory=trials_directory)
model_class = link_bot_classifiers.get_model(params['model_class'])
###############
# Dataset
###############
dataset = ClassifierDatasetLoader(
dataset_dirs,
load_true_states=True,
use_gt_rope=use_gt_rope,
threshold=params['classifier_dataset_hparams']['labeling_params']
['threshold'],
old_compat=old_compat)
model = model_class(hparams=params,
batch_size=batch_size,
scenario=dataset.scenario)
tf_dataset = dataset.get_datasets(mode=mode)
scenario = dataset.scenario
###############
# Evaluate
###############
tf_dataset = batch_tf_dataset(tf_dataset, batch_size, drop_remainder=True)
model = classifier_utils.load_generic_model([checkpoint])
for batch_idx, example in enumerate(
progressbar(tf_dataset, widgets=base_dataset.widgets)):
example.update(dataset.batch_metadata)
predictions, _ = model.check_constraint_from_example(example,
training=False)
labels = tf.expand_dims(example['is_close'][:, 1:], axis=2)
probabilities = predictions['probabilities']
# Visualization
example.pop("time")
example.pop("batch_size")
decisions = probabilities > 0.5
classifier_is_correct = tf.squeeze(tf.equal(decisions,
tf.cast(labels, tf.bool)),
axis=-1)
for b in range(batch_size):
example_b = index_dict_of_batched_tensors_tf(example, b)
# if the classifier is correct at all time steps, ignore
if only_errors and tf.reduce_all(classifier_is_correct[b]):
continue
def _custom_viz_t(scenario: Base3DScenario, e: Dict, t: int):
if t > 0:
accept_probability_t = predictions['probabilities'][
b, t - 1, 0].numpy()
else:
accept_probability_t = -999
scenario.plot_accept_probability(accept_probability_t)
traj_idx_msg = Float32()
traj_idx_msg.data = batch_idx * batch_size + b
traj_idx_pub_.publish(traj_idx_msg)
anim = RvizAnimation(scenario=scenario,
n_time_steps=dataset.horizon,
init_funcs=[
init_viz_env,
dataset.init_viz_action(),
],
t_funcs=[
_custom_viz_t,
dataset.classifier_transition_viz_t(),
ExperimentScenario.plot_stdev_t,
])
with open("debugging.hjson", 'w') as f:
my_hdump(numpify(example_b), f)
anim.play(example_b)
def filter(self, environment: Dict, state: Optional[Dict], observation: Dict) -> Tuple[Dict, Dict]:
mean_state, stdev_state = self.filter_differentiable(environment=environment, state=state, observation=observation)
return numpify(mean_state), numpify(stdev_state)
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 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)
