visualize_1_rail_3_joint_robot.py

import numpy as np
from matplotlib import pyplot as plt
from scipy.stats import norm

from src.inverse_kinematics_sbi.plot import plot_arms
from src.inverse_kinematics_sbi.proba_robot import ProbaRobot
from src.inverse_kinematics_sbi.robots import Joint, Rail, RobotArm
from src.inverse_kinematics_sbi.solve_abc import distance_end_effector_position, draw_abc_samples

lengths = [0.5, 0.5, 1.0]
sigmas = [0.25, 0.5, 0.5, 0.5]

robot = RobotArm(components=[Rail(np.pi / 2)] + [Joint(length=length) for length in lengths])
prior = norm(np.zeros_like(sigmas), sigmas)
prob_robot = ProbaRobot(robot, prior)
observations = np.array([[1.8, 0.0]])
n_sample = 50
target = observations[0]
parameters = prob_robot.sample_posterior(observations, n_sample=n_sample, verbose=True)
jacobians = robot.forward_jacobian(parameters)
post_prob = (
    np.prod(prior.pdf(parameters), axis=2)
    * np.linalg.det(np.matmul(jacobians, jacobians.transpose((0, 1, 3, 2)))) ** (-1 / 2)
)[0]
sorted_indices = np.argsort(post_prob)
coords = robot.forward(parameters, return_intermediates=True)[0, sorted_indices, 1:]
# parameters = draw_abc_samples(n_sample, robot.forward, lambda n_: prior.rvs((n_, robot.get_n_params())), distance_end_effector_position, target, 0.01, verbose=True)
plt.figure()
plt.xlim([0, 2])
plt.ylim([-1, 1])
plot_arms(coords, None, density=post_prob[sorted_indices])