256 lines
9.1 KiB
Python
256 lines
9.1 KiB
Python
import pinocchio as pin
|
|
import numpy as np
|
|
try:
|
|
import meshcat
|
|
import meshcat.geometry as g
|
|
import meshcat.transformations as tf
|
|
except:
|
|
print("You need to install meshcat.")
|
|
|
|
import crocoddyl
|
|
|
|
|
|
|
|
class ResidualForce3D(crocoddyl.ResidualModelAbstract):
|
|
def __init__(self, state, contact_name, nu):
|
|
crocoddyl.ResidualModelAbstract.__init__(self, state, 3, nu, True, True, True)
|
|
self.contact_name = contact_name
|
|
|
|
def calc(self, data, x, u=None):
|
|
data.r = data.shared.contacts.contacts[self.contact_name].f.vector[:3]
|
|
|
|
def calcDiff(self, data, x, u=None):
|
|
data.Rx = data.shared.contacts.contacts[self.contact_name].df_dx[:3]
|
|
data.Ru = data.shared.contacts.contacts[self.contact_name].df_du[:3]
|
|
|
|
|
|
class ResidualFrictionCone(crocoddyl.ResidualModelAbstract):
|
|
def __init__(self, state, contact_name, mu, nu):
|
|
crocoddyl.ResidualModelAbstract.__init__(self, state, 1, nu, True, True, True)
|
|
self.mu = mu
|
|
self.contact_name = contact_name
|
|
|
|
self.dcone_df = np.zeros((1, 3))
|
|
self.df_dx = np.zeros((3, self.state.ndx))
|
|
self.df_du = np.zeros((3, self.nu))
|
|
|
|
def calc(self, data, x, u=None):
|
|
F = data.shared.contacts.contacts[self.contact_name].f.vector[:3]
|
|
data.r[0] = np.array([self.mu * F[2] - np.sqrt(F[0]**2 + F[1]**2)])
|
|
|
|
def calcDiff(self, data, x, u=None):
|
|
F = data.shared.contacts.contacts[self.contact_name].f.vector[:3]
|
|
|
|
self.dcone_df[0, 0] = -F[0] / np.sqrt(F[0]**2 + F[1]**2)
|
|
self.dcone_df[0, 1] = -F[1] / np.sqrt(F[0]**2 + F[1]**2)
|
|
self.dcone_df[0, 2] = self.mu
|
|
|
|
|
|
self.df_dx = data.shared.contacts.contacts[self.contact_name].df_dx[:3]
|
|
self.df_du = data.shared.contacts.contacts[self.contact_name].df_du[:3]
|
|
|
|
data.Rx = self.dcone_df @ self.df_dx
|
|
data.Ru = self.dcone_df @ self.df_du
|
|
|
|
|
|
def meshcat_material(r, g, b, a):
|
|
material = meshcat.geometry.MeshPhongMaterial()
|
|
material.color = int(r * 255) * 256 ** 2 + int(g * 255) * 256 + int(b * 255)
|
|
material.opacity = a
|
|
material.linewidth = 5.0
|
|
return material
|
|
|
|
def addViewerBox(viz, name, sizex, sizey, sizez, rgba):
|
|
if isinstance(viz, pin.visualize.MeshcatVisualizer):
|
|
viz.viewer[name].set_object(meshcat.geometry.Box([sizex, sizey, sizez]),
|
|
meshcat_material(*rgba))
|
|
|
|
|
|
def addLineSegment(viz, name, vertices, rgba):
|
|
if isinstance(viz, pin.visualize.MeshcatVisualizer):
|
|
viz.viewer[name].set_object(meshcat.geometry.LineSegments(
|
|
meshcat.geometry.PointsGeometry(np.array(vertices)),
|
|
meshcat_material(*rgba)
|
|
))
|
|
|
|
def addPoint(viz, name, vertices, rgba):
|
|
if isinstance(viz, pin.visualize.MeshcatVisualizer):
|
|
viz.viewer[name].set_object(meshcat.geometry.Points(
|
|
meshcat.geometry.PointsGeometry(np.array(vertices)),
|
|
meshcat_material(*rgba)
|
|
))
|
|
|
|
def meshcat_transform(x, y, z, q, u, a, t):
|
|
return np.array(pin.XYZQUATToSE3([x, y, z, q, u, a, t]))
|
|
|
|
def applyViewerConfiguration(viz, name, xyzquat):
|
|
if isinstance(viz, pin.visualize.MeshcatVisualizer):
|
|
viz.viewer[name].set_transform(meshcat_transform(*xyzquat))
|
|
|
|
def get_solution_trajectories(solver, rmodel, rdata, supportFeetIds):
|
|
xs, us = solver.xs, solver.us
|
|
nq, nv, N = rmodel.nq, rmodel.nv, len(xs)
|
|
jointPos_sol = []
|
|
jointVel_sol = []
|
|
jointAcc_sol = []
|
|
jointTorques_sol = []
|
|
centroidal_sol = []
|
|
|
|
x = []
|
|
for time_idx in range (N):
|
|
q, v = xs[time_idx][:nq], xs[time_idx][nq:]
|
|
pin.framesForwardKinematics(rmodel, rdata, q)
|
|
pin.computeCentroidalMomentum(rmodel, rdata, q, v)
|
|
centroidal_sol += [
|
|
np.concatenate(
|
|
[pin.centerOfMass(rmodel, rdata, q, v), np.array(rdata.hg)]
|
|
)
|
|
]
|
|
jointPos_sol += [q]
|
|
jointVel_sol += [v]
|
|
x += [xs[time_idx]]
|
|
if time_idx < N-1:
|
|
jointAcc_sol += [solver.problem.runningDatas[time_idx].xnext[nq::]]
|
|
jointTorques_sol += [us[time_idx]]
|
|
|
|
|
|
|
|
|
|
sol = {'x':x, 'centroidal':centroidal_sol, 'jointPos':jointPos_sol,
|
|
'jointVel':jointVel_sol, 'jointAcc':jointAcc_sol,
|
|
'jointTorques':jointTorques_sol}
|
|
|
|
|
|
for frame_idx in supportFeetIds:
|
|
# print('extract foot id ', frame_idx, "_name = ", rmodel.frames[frame_idx].name)
|
|
ct_frame_name = rmodel.frames[frame_idx].name + "_contact"
|
|
datas = [solver.problem.runningDatas[i].differential.multibody.contacts.contacts[ct_frame_name] for i in range(N-1)]
|
|
ee_forces = [datas[k].f.vector for k in range(N-1)]
|
|
sol[ct_frame_name] = [ee_forces[i] for i in range(N-1)]
|
|
|
|
return sol
|
|
|
|
|
|
|
|
|
|
class Arrow(object):
|
|
def __init__(self, meshcat_vis, name,
|
|
location=[0,0,0],
|
|
vector=[0,0,1],
|
|
length_scale=1,
|
|
color=0xff0000):
|
|
|
|
self.vis = meshcat_vis[name]
|
|
self.cone = self.vis["cone"]
|
|
self.line = self.vis["line"]
|
|
self.material = g.MeshBasicMaterial(color=color, reflectivity=0.5)
|
|
|
|
self.location, self.length_scale = location, length_scale
|
|
self.anchor_as_vector(location, vector)
|
|
|
|
def _update(self):
|
|
# pass
|
|
translation = tf.translation_matrix(self.location)
|
|
rotation = self.orientation
|
|
offset = tf.translation_matrix([0, self.length/2, 0])
|
|
self.pose = translation @ rotation @ offset
|
|
self.vis.set_transform(self.pose)
|
|
|
|
def set_length(self, length, update=True):
|
|
self.length = length * self.length_scale
|
|
cone_scale = self.length/0.08
|
|
self.line.set_object(g.Cylinder(height=self.length, radius=0.005), self.material)
|
|
self.cone.set_object(g.Cylinder(height=0.015,
|
|
radius=0.01,
|
|
radiusTop=0.,
|
|
radiusBottom=0.01),
|
|
self.material)
|
|
self.cone.set_transform(tf.translation_matrix([0.,cone_scale*0.04,0]))
|
|
if update:
|
|
self._update()
|
|
|
|
def set_direction(self, direction, update=True):
|
|
orientation = np.eye(4)
|
|
orientation[:3, 0] = np.cross([1,0,0], direction)
|
|
orientation[:3, 1] = direction
|
|
orientation[:3, 2] = np.cross(orientation[:3, 0], orientation[:3, 1])
|
|
self.orientation = orientation
|
|
if update:
|
|
self._update()
|
|
|
|
def set_location(self, location, update=True):
|
|
self.location = location
|
|
if update:
|
|
self._update()
|
|
|
|
def anchor_as_vector(self, location, vector, update=True):
|
|
self.set_direction(np.array(vector)/np.linalg.norm(vector), False)
|
|
self.set_location(location, False)
|
|
self.set_length(np.linalg.norm(vector), False)
|
|
if update:
|
|
self._update()
|
|
|
|
def delete(self):
|
|
self.vis.delete()
|
|
|
|
|
|
|
|
|
|
class Cone(object):
|
|
def __init__(self, meshcat_vis, name,
|
|
location=[0,0,0], mu=1,
|
|
vector=[0,0,1],
|
|
length_scale=0.06):
|
|
|
|
self.vis = meshcat_vis[name]
|
|
self.cone = self.vis["cone"]
|
|
self.material = g.MeshBasicMaterial(color=0xffffff, opacity = 0.5, reflectivity=0.5)
|
|
|
|
|
|
self.mu = mu * length_scale
|
|
self.location, self.length_scale = location, length_scale
|
|
self.anchor_as_vector(location, vector)
|
|
|
|
def _update(self):
|
|
# pass
|
|
translation = tf.translation_matrix(self.location)
|
|
rotation = self.orientation
|
|
offset = tf.translation_matrix([0, self.length/2, 0])
|
|
self.pose = translation @ rotation @ offset
|
|
self.vis.set_transform(self.pose)
|
|
|
|
def set_length(self, length, update=True):
|
|
self.length = length * self.length_scale
|
|
cone_scale = self.length
|
|
self.cone.set_object(g.Cylinder(height=cone_scale,
|
|
radius=self.mu,
|
|
radiusTop=self.mu,
|
|
radiusBottom=0),
|
|
self.material)
|
|
# self.cone.set_transform(tf.translation_matrix([0.,cone_scale*0.04,0]))
|
|
if update:
|
|
self._update()
|
|
|
|
def set_direction(self, direction, update=True):
|
|
orientation = np.eye(4)
|
|
orientation[:3, 0] = np.cross([1,0,0], direction)
|
|
orientation[:3, 1] = direction
|
|
orientation[:3, 2] = np.cross(orientation[:3, 0], orientation[:3, 1])
|
|
self.orientation = orientation
|
|
if update:
|
|
self._update()
|
|
|
|
def set_location(self, location, update=True):
|
|
self.location = location
|
|
if update:
|
|
self._update()
|
|
|
|
def anchor_as_vector(self, location, vector, update=True):
|
|
self.set_direction(np.array(vector)/np.linalg.norm(vector), False)
|
|
self.set_location(location, False)
|
|
self.set_length(np.linalg.norm(vector), False)
|
|
if update:
|
|
self._update()
|
|
|
|
def delete(self):
|
|
self.vis.delete() |