# Imports
import numpy as np
import altair as alt
import pydot
from IPython.display import display, SVG, clear_output
import plotly.express as px
import time
from collections import namedtuple
import matplotlib.pyplot as plt
from pydrake.all import (
AddMultibodyPlantSceneGraph, AngleAxis, BasicVector,
DiagramBuilder, InverseKinematics, FindResourceOrThrow, Integrator, JacobianWrtVariable,
LeafSystem, MeshcatVisualizerCpp, MultibodyPlant, MultibodyPositionToGeometryPose, Parser,
PiecewisePose, PiecewisePolynomial, Quaternion, RigidTransform, Solve, MeshcatVisualizerParams,
RollPitchYaw, RotationMatrix, SceneGraph, Simulator, TrajectorySource, InverseDynamicsController
)
from pydrake.examples.manipulation_station import ManipulationStation
from manipulation import running_as_notebook
from manipulation.meshcat_cpp_utils import (
StartMeshcat, MeshcatJointSlidersThatPublish)
from manipulation.scenarios import AddMultibodyTriad, SetColor
import pandas as pd
# Start the visualizer.
meshcat = StartMeshcat()
def ApplyIK(plant, plant_context, the_pose, ee_angle, q0, uarm):
ik = InverseKinematics(plant, plant_context)
#ik.AddPositionConstraint(
# plant.GetFrameByName("uarm_link_ee", uarm), [0.0, 0.0, 0.0], plant.GetFrameByName("root"),
# the_pose.translation(), the_pose.translation())
ik.AddOrientationConstraint(
plant.GetFrameByName("uarm_link_ee", uarm), RotationMatrix(), plant.GetFrameByName("root"),
ee_angle.rotation(), 0.0)
ik.AddMinimumDistanceConstraint(0.001, 0.1)
prog = ik.get_mutable_prog()
q = ik.q()
prog.AddQuadraticErrorCost(np.identity(len(q)), q0, q)
prog.SetInitialGuess(q, q0)
result = Solve(ik.prog())
#clear_output(wait=True)
if result.is_success():
placeholder=0# print("IK success")
else:
print("IK failure")
def uarm_sim():
builder = DiagramBuilder()
plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)
parser = Parser(plant, scene_graph)
uarm = parser.AddModelFromFile("4DOF_Robot.urdf")
brick = parser.AddModelFromFile(FindResourceOrThrow(
"drake/examples/manipulation_station/models/061_foam_brick.sdf"))
plant.Finalize()
visualizer = MeshcatVisualizerCpp.AddToBuilder(
builder,
scene_graph,
meshcat,
MeshcatVisualizerParams(delete_prefix_initialization_event=False))
diagram = builder.Build()
context = diagram.CreateDefaultContext()
plant_context = plant.GetMyContextFromRoot(context)
B_O = plant.GetBodyByName("base_link", brick)
uarm_base = plant.GetBodyByName("root")
uarm_ee = plant.GetBodyByName("uarm_link_ee")
joint_0 = plant.GetJointByName("uarm_joint_0")
joint_1 = plant.GetJointByName("uarm_joint_1")
joint_2 = plant.GetJointByName("uarm_joint_2")
joint_3 = plant.GetJointByName("uarm_joint_3")
joint_4 = plant.GetJointByName("uarm_joint_4")
X_WO = plant.EvalBodyPoseInWorld(plant_context, uarm_base)
p_GgraspO = [0.14, 0, -0.085]
R_GgraspO = RotationMatrix.MakeXRotation(np.pi*2.0).multiply(
RotationMatrix.MakeZRotation(np.pi*2.0))
X_GgraspO = RigidTransform(R_GgraspO, p_GgraspO)
X_OGgrasp = X_GgraspO.inverse()
X_WGgrasp = X_WO.multiply(X_OGgrasp)
plant.SetFreeBodyPose(plant_context, uarm_base, X_WGgrasp)
plant.GetJointByName("uarm_joint_0", uarm).set_angle(plant_context, 0)
plant.GetJointByName("uarm_joint_1", uarm).set_angle(plant_context, 0)
plant.GetJointByName("uarm_joint_2", uarm).set_angle(plant_context, 0)
plant.GetJointByName("uarm_joint_3", uarm).set_angle(plant_context, 0)
initialpos = plant.CalcRelativeTransform(plant_context, plant.GetFrameByName("root"), plant.GetFrameByName("uarm_link_ee", uarm))
"""
#DEGBUGGING/Future Use
plant.GetJointByName("uarm_joint_0", uarm).set_angle(plant_context, np.pi)
plant.GetJointByName("uarm_joint_1", uarm).set_angle(plant_context, np.pi/4)
plant.GetJointByName("uarm_joint_2", uarm).set_angle(plant_context, -np.pi/2)
plant.GetJointByName("uarm_joint_3", uarm).set_angle(plant_context, np.pi/4)
initial_grip_pos = plant.CalcRelativeTransform(plant_context, plant.GetFrameByName("root"), plant.GetFrameByName("uarm_link_ee", uarm))
plant.GetJointByName("uarm_joint_0", uarm).set_angle(plant_context, np.pi/2)
plant.GetJointByName("uarm_joint_1", uarm).set_angle(plant_context, np.pi/4)
plant.GetJointByName("uarm_joint_2", uarm).set_angle(plant_context, -np.pi/4)
plant.GetJointByName("uarm_joint_3", uarm).set_angle(plant_context, 0)
inter_pos = plant.CalcRelativeTransform(plant_context, plant.GetFrameByName("root"), plant.GetFrameByName("uarm_link_ee", uarm))
plant.GetJointByName("uarm_joint_0", uarm).set_angle(plant_context, 0)
plant.GetJointByName("uarm_joint_1", uarm).set_angle(plant_context, np.pi/4)
plant.GetJointByName("uarm_joint_2", uarm).set_angle(plant_context, -np.pi/2)
plant.GetJointByName("uarm_joint_3", uarm).set_angle(plant_context, np.pi/4)
finalpos = plant.CalcRelativeTransform(plant_context, plant.GetFrameByName("root"), plant.GetFrameByName("uarm_link_ee", uarm))
"""
time.sleep(6)
positions = initialpos
dt = 0.01
t_vec = np.arange(0,10,dt)
simulation_size = len(t_vec) # We use this alot, so lets store it.
joint0angles = np.zeros(2*simulation_size)
joint1angles = np.zeros(2*simulation_size)
joint2angles = np.zeros(2*simulation_size)
for i in range(1, len(t_vec)):
q0 = plant.GetPositions(plant_context)
ApplyIK(plant, plant_context, positions, initialpos, q0, uarm)
diagram.Publish(context)
#x_val = 0.14*np.cos(t_vec[i]*(np.pi/10))
#y_val = 0.1*np.sin(t_vec[i]*(np.pi/10))
#z_val = 0.1*np.sin(t_vec[i]*(np.pi/10)) -0.085
j_0 = np.pi*np.sin(t_vec[i]*(np.pi/20))
j_1 = (np.pi/4)*(np.sin(t_vec[i]*(np.pi/20.0))+np.sin(t_vec[i]*(np.pi/10.0)))
j_2 = -(np.pi/2)*np.sin(t_vec[i]*(np.pi/20))
joint0angles[i] = j_0
joint1angles[i] = j_1
joint2angles[i] = j_2
plant.GetJointByName("uarm_joint_0", uarm).set_angle(plant_context, j_0)
plant.GetJointByName("uarm_joint_1", uarm).set_angle(plant_context, j_1)
plant.GetJointByName("uarm_joint_2", uarm).set_angle(plant_context, j_2)
positions = plant.CalcRelativeTransform(plant_context, plant.GetFrameByName("root"), plant.GetFrameByName("uarm_link_ee", uarm))
joint0angles[simulation_size] = j_0
joint1angles[simulation_size] = j_1
joint2angles[simulation_size] = j_2
time.sleep(2)
for i in range(1, len(t_vec)):
q0 = plant.GetPositions(plant_context)
ApplyIK(plant, plant_context, positions, initialpos, q0, uarm)
diagram.Publish(context)
#x_val = 0.14*np.cos(t_vec[i]*(np.pi/10))
#y_val = 0.1*np.sin(t_vec[i]*(np.pi/10))
#z_val = 0.1*np.sin(t_vec[i]*(np.pi/10)) -0.085
j_0 = (np.pi/2)*np.sin(t_vec[i]*(np.pi/20)) + np.pi
j_1 = np.pi/4*(np.cos(t_vec[i]*(np.pi/20.0))+np.pi/4*np.sin(t_vec[i]*(np.pi/10.0)))
j_2 = (np.pi/2)*np.sin(t_vec[i]*(np.pi/20)) - np.pi/2
joint0angles[simulation_size+i] = j_0
joint1angles[simulation_size+i] = j_1
joint2angles[simulation_size+i] = j_2
plant.GetJointByName("uarm_joint_0", uarm).set_angle(plant_context, j_0)
plant.GetJointByName("uarm_joint_1", uarm).set_angle(plant_context, j_1)
plant.GetJointByName("uarm_joint_2", uarm).set_angle(plant_context, j_2)
X_WO = plant.EvalBodyPoseInWorld(plant_context, uarm_ee)
plant.SetFreeBodyPose(plant_context, B_O, X_WO.multiply(RigidTransform(R_GgraspO, [0, 0, -0.085])))
positions = plant.CalcRelativeTransform(plant_context, plant.GetFrameByName("root"), plant.GetFrameByName("uarm_link_ee", uarm))
tsim_vec = np.arange(0,20,dt)
plt.plot(tsim_vec,joint0angles, tsim_vec,joint1angles, tsim_vec,joint2angles)
plt.legend(['joint0', 'joint1', 'joint2'])
plt.grid()
plt.xlabel("Time (s)")
plt.ylabel("Joint Angle (rad)")
plt.show()
uarm_sim()