Source code for gmmvi.experiments.target_distributions.planar_robot

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp

from gmmvi.experiments.evaluation.visualize_planar_robot import visualize_mixture
from gmmvi.experiments.target_distributions.lnpdf import LNPDF
from gmmvi.models.gmm_wrapper import GmmWrapper

tfd = tfp.distributions

import matplotlib.pyplot as plt

[docs]class PlanarRobot(LNPDF): """This class reimplements the "PlanarRobot" experiments used by :cite:t:`Arenz2020`. Parameters: num_links: int The number of links of the robot num_goals: int The number of goal positions, must be either 1 or 4 prior_std: float The standard deviation of the (zero-mean) prior on the joint angles. The first value is ignored, as the first link always has a standard deviation of 1. likelihood_std: float The std-deviation used for penalizing the distance in X-Y between the robot endeffective and the goal position. """ def __init__(self, num_links, num_goals, prior_std=2e-1, likelihood_std=1e-2): super(PlanarRobot, self).__init__(use_log_density_and_grad=False) self._num_dimensions = num_links prior_stds = prior_std * np.ones(num_links) prior_stds[0] = 1. self.prior = tfd.MultivariateNormalDiag(loc=tf.zeros(num_links), scale_diag=prior_stds.astype(np.float32)) self.link_lengths = np.ones(self._num_dimensions) self._num_goals = num_goals if num_goals == 1: self.goals = tf.constant([[7., 0.]], dtype=tf.float32) elif num_goals == 4: self.goals = tf.constant([[7., 0.], [-7., 0.], [0., 7.], [0., -7.]], dtype=tf.float32) else: raise ValueError self.goal_Gaussians = [] for goal in self.goals: self.goal_Gaussians.append(tfd.MultivariateNormalDiag(loc=goal, scale_identity_multiplier=likelihood_std))
[docs] def likelihood(self, pos: tf.Tensor) -> tf.Tensor: likelihoods = tf.TensorArray(size=self._num_goals, dtype=tf.float32) for i in range(self._num_goals): likelihoods = likelihoods.write(i, self.goal_Gaussians[i].log_prob(pos)) return tf.reduce_max(likelihoods.stack(), axis=0)
[docs] def get_num_dimensions(self): return self._num_dimensions
[docs] def forward_kinematics(self, theta): y = tf.zeros(tf.shape(theta)[0]) x = tf.zeros(tf.shape(theta)[0]) for i in range(0, self._num_dimensions): y += self.link_lengths[i] * tf.math.sin(tf.reduce_sum(theta[:, :i + 1], axis=1)) x += self.link_lengths[i] * tf.math.cos(tf.reduce_sum(theta[:, :i + 1], axis=1)) return tf.stack((x, y), axis=1)
[docs] def log_density(self, theta): return self.prior.log_prob(theta) + self.likelihood(self.forward_kinematics(theta))
[docs] def expensive_metrics(self, model: GmmWrapper, samples: tf.Tensor) -> dict: """ This method computes two task-specific metrics: 1. The number of detected modes: This is course heuristic to count the different configurations used for reaching each of the goal positions (potentially misleading!) 2. Plots of the mean configurations of the learned model Parameters: model: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The learned model that we want to evaluate for this target distribution. samples: tf.Tensor Samples that have been drawn from the model, which can be used for evaluations. Returns: dict: a dictionary containing two items (the number of detected modes, and a figure showing the mean configurations). """ all_comp_colors = np.repeat("k", model.num_components) colors = ["b", "g", "r", "c", "m", "y"] next_color_idx = 0 expensive_metrics = dict() for goal in self.goals: end_eff_error = tf.norm( self.forward_kinematics(model.means) - goal, axis=1) close_to_goal = (end_eff_error < 0.05) good_components = self.log_density(model.means) > -7. this_goal_indices = tf.reshape(tf.where(tf.reduce_all([close_to_goal, good_components], axis=0)), [-1]) if tf.size(this_goal_indices) == 0: num_detected_modes = 0 else: this_goal_first_links = tf.gather(model.means[:, 0], this_goal_indices) try: first_angles_sorted_ind = tf.argsort(this_goal_first_links) except: print("debug") first_angles_sorted = tf.gather(this_goal_first_links, first_angles_sorted_ind) gaps = first_angles_sorted[1:] - first_angles_sorted[:-1] comp_colors = np.repeat(colors[next_color_idx], tf.size(first_angles_sorted_ind)) num_detected_modes = 1 clusters = [] splits = tf.reshape(tf.where(gaps > 0.4), [-1]) start = 0 for i in range(len(splits)): clusters.append(np.mean(first_angles_sorted[start:splits[i] + 1])) start = splits[i] + 1 num_detected_modes += 1 next_color_idx += 1 next_color_idx = next_color_idx % len(colors) next_color = colors[next_color_idx] comp_colors[splits[i]:] = next_color all_comp_colors[this_goal_indices.numpy()[first_angles_sorted_ind.numpy()]] = comp_colors clusters.append(np.mean(first_angles_sorted[start:])) print(f"clusters for goal at position [{goal[0]}, {goal[1]}]: {clusters}") expensive_metrics.update({f"num_detected_modes_{goal}": num_detected_modes}) mixture_fig = plt.figure(1) visualize_mixture(model.weights.numpy(), model.means.numpy()) #, comp_colors=all_comp_colors) mixture_fig.tight_layout() expensive_metrics.update( {"mixture_plot": mixture_fig}) return expensive_metrics
def make_single_goal(): return PlanarRobot(10, 1) def make_four_goal(): return PlanarRobot(10, 4)