def __call__(self, environment: Dict, state: Dict):
# sample a bunch of actions (batched?) and pick the best one
max_unstuck_probability = -1
best_action = None
# anim = RvizAnimationController(np.arange(self.n_action_samples))
# while not anim.done:
for _ in range(self.n_action_samples):
self.scenario.last_action = None
action = self.scenario.sample_action(
action_rng=self.action_rng,
environment=environment,
state=state,
action_params=self.data_collection_params,
validate=True)
# TODO: use the unconstrained dynamics to predict the state resulting from (e, s, a)
# then add that to the recovery_model_input
recovery_model_input = environment
recovery_model_input.update(
add_batch(state)) # add time dimension to state and action
recovery_model_input.update(add_batch(action))
recovery_model_input = make_dict_tf_float32(
add_batch(recovery_model_input))
recovery_model_input.update({
'batch_size': 1,
'time': 2,
})
recovery_model_output = self.model(recovery_model_input,
training=False)
recovery_probability = recovery_model_output['probabilities']
# self.scenario.plot_environment_rviz(environment)
# self.scenario.plot_state_rviz(state, label='stuck state')
# self.scenario.plot_recovery_probability(recovery_probability)
# color_factor = log_scale_0_to_1(tf.squeeze(recovery_probability), k=100)
# self.scenario.plot_action_rviz(state, action, label='proposed', color=cm.Greens(color_factor), idx=1)
if recovery_probability > max_unstuck_probability:
max_unstuck_probability = recovery_probability
best_action = action
# print(max_unstuck_probability)
# self.scenario.plot_action_rviz(state, action, label='best_proposed', color='g', idx=2)
# anim.step()
best_action_noisy = self.scenario.add_noise(best_action,
self.noise_rng)
return best_action_noisy
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 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 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 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 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 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 _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
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 add_time(self):
self.dnames.insert(1, "time")
self.data = add_batch(self.data)
return self
def remove_batch(self):
self.dnames.remove("batch")
self.data = add_batch(self.data)
return self
def add_batch(self):
self.dnames.insert(0, "batch")
self.data = add_batch(self.data)
return self