as we already started discussing, the code could probably use some refactoring because at the moment it's a bit messy. I haven't had the time to think about it properly, but I have recently developed a simple simulator in Python for a class I am teaching. I've tried to make it as simple as possible so that student could understand it easily. I'm sharing it with you here. Maybe this code could be a starting point to think of the new code
import pinocchio as se3
import numpy as np
from numpy.linalg import norm
import os
import math
import time
class ContactPoint:
''' A point on the robot surface that can make contact with surfaces.
'''
def __init__(self, model, data, frame_name):
self.model = model # robot model
self.data = data # robot data
self.frame_name = frame_name # name of reference frame associated to this contact point
self.frame_id = model.getFrameId(frame_name) # id of the reference frame
self.active = False # True if this contact point is in contact
def get_position(self):
''' Get the current position of this contact point
'''
M = self.data.oMf[self.frame_id]
return M.translation
def get_velocity(self):
M = self.data.oMf[self.frame_id]
R = se3.SE3(M.rotation, 0*M.translation) # same as M but with translation set to zero
v_local = se3.getFrameVelocity(self.model, self.data, self.frame_id)
v_world = (R.act(v_local)).linear # convert velocity from local frame to world frame
return v_world
def get_jacobian(self):
J6 = se3.getFrameJacobian(self.model, self.data, self.frame_id, se3.ReferenceFrame.LOCAL_WORLD_ALIGNED)
return J6[:3,:]
class ContactSurface:
''' A visco-elastic planar surface
'''
def __init__(self, name, pos, normal, K, B, mu):
self.name = name # name of this contact surface
self.x0 = pos # position of a point of the surface
self.normal = normal # direction of the normal to the surface
self.K = K # stiffness of the surface material
self.B = B # damping of the surface material
self.mu = mu # friction coefficient of the surface
self.bias = self.x0.dot(self.normal)
def check_collision(self, p):
''' Check the collision of the given point
with this contact surface. If the point is not
inside this surface, then return False.
'''
normal_penetration = self.bias - p.dot(self.normal)
if(normal_penetration < 0.0):
return False # no penetration
return True
def compute_force(self, contact_point, anchor_point):
cp = contact_point
p0 = anchor_point
p = cp.get_position()
v = cp.get_velocity()
# compute contact force using spring-damper law
f = self.K.dot(p0 - p) - self.B.dot(v)
# check whether contact force is outside friction cone
f_N = f.dot(self.normal) # norm of normal force
f_T = f - f_N*self.normal # tangential force (3d)
f_T_norm = norm(f_T) # norm of tangential force
if(f_T_norm > self.mu*f_N):
# contact is slipping
t_dir = f_T / f_T_norm # direction of tangential force
# saturate force at the friction cone boundary
f = f_N*self.normal + self.mu*f_N*t_dir
# update anchor point so that f is inside friction cone
delta_p0 = (f_T_norm - self.mu*f_N) / self.K[0,0]
p0 -= t_dir*delta_p0
return f, p0
class Contact:
''' A contact between a contact-point and a contact-surface
'''
def __init__(self, contact_point, contact_surface):
self.cp = contact_point
self.cs = contact_surface
self.reset_contact_position()
def reset_contact_position(self):
# Initialize anchor point p0, that is the initial (0-load) position of the spring
self.p0 = self.cp.get_position()
self.in_contact = True
def compute_force(self):
self.f, self.p0 = self.cs.compute_force(self.cp, self.p0)
return self.f
def get_jacobian(self):
return self.cp.get_jacobian()
class RobotSimulator:
# Class constructor
def __init__(self, conf, robot):
self.conf = conf
self.robot = robot
self.model = self.robot.model
self.data = self.model.createData()
self.t = 0.0 # time
self.nv = nv = self.model.nv # Dimension of joint velocities vector
self.na = na = robot.na # number of actuated joints
# Matrix S used as filter of vetor of inputs U
self.S = np.hstack((np.zeros((na, nv-na)), np.eye(na, na)))
self.contacts = []
self.candidate_contact_points = [] # candidate contact points
self.contact_surfaces = []
self.init(conf.q0, None, True)
self.tau_c = np.zeros(na) # Coulomb friction torque
self.simulate_coulomb_friction = conf.simulate_coulomb_friction
self.simulation_type = conf.simulation_type
if(self.simulate_coulomb_friction):
self.tau_coulomb_max = 1e-2*conf.tau_coulomb_max*self.model.effortLimit
else:
self.tau_coulomb_max = np.zeros(na)
# Re-initialize the simulator
def init(self, q0=None, v0=None, reset_contact_positions=False):
self.first_iter = True
if q0 is not None:
self.q = q0.copy()
if(v0 is None):
self.v = np.zeros(self.robot.nv)
else:
self.v = v0.copy()
self.dv = np.zeros(self.robot.nv)
self.resize_contact_data(reset_contact_positions)
def resize_contact_data(self, reset_contact_positions=False):
self.nc = len(self.contacts) # number of contacts
self.nk = 3*self.nc # size of contact force vector
self.f = np.zeros(self.nk) # contact forces
self.Jc = np.zeros((self.nk, self.model.nv)) # contact Jacobian
# reset contact position
if(reset_contact_positions):
se3.forwardKinematics(self.model, self.data, self.q)
se3.updateFramePlacements(self.model, self.data)
for c in self.contacts:
c.reset_contact_position()
self.compute_forces(compute_data=True)
def add_candidate_contact_point(self, frame_name):
self.candidate_contact_points += [ContactPoint(self.model, self.data, frame_name)]
def add_contact_surface(self, name, pos, normal, K, B, mu):
''' Add a contact surface (i.e., a wall) located at "pos", with normal
outgoing direction "normal", 3d stiffness K, 3d damping B.
'''
self.contact_surfaces += [ContactSurface(name, pos, normal, K, B, mu)]
def collision_detection(self):
for s in self.contact_surfaces: # for each contact surface
for cp in self.candidate_contact_points: # for each candidate contact point
p = cp.get_position()
if(s.check_collision(p)): # check whether the point is colliding with the surface
if(not cp.active): # if the contact was not already active
print("Collision detected between point", cp.frame_name, " at ", p)
cp.active = True
cp.contact = Contact(cp, s)
self.contacts += [cp.contact]
self.resize_contact_data()
else:
if(cp.active):
print("Contact lost between point", cp.frame_name, " at ", p)
cp.active = False
self.contacts.remove(cp.contact)
self.resize_contact_data()
def compute_forces(self, compute_data=True):
'''Compute the contact forces from q, v and elastic model'''
if compute_data:
se3.forwardKinematics(self.model, self.data, self.q, self.v)
# Computes the placements of all the operational frames according to the current joint placement stored in data
se3.updateFramePlacements(self.model, self.data)
self.collision_detection()
i = 0
for c in self.contacts:
self.f[i:i+3] = c.compute_force()
self.Jc[i:i+3, :] = c.get_jacobian()
i += 3
return self.f
def step(self, u, dt):
# (Forces are directly in the world frame, and aba wants them in the end effector frame)
se3.computeAllTerms(self.model, self.data, self.q, self.v)
se3.updateFramePlacements(self.model, self.data)
M = self.data.M
h = self.data.nle
self.collision_detection()
self.compute_forces(False)
self.tau_c = self.tau_coulomb_max*np.sign(self.v[-self.na:])
self.dv = np.linalg.solve(M, self.S.T.dot(u-self.tau_c) - h + self.Jc.T.dot(self.f))
v_mean = self.v + 0.5*dt*self.dv
self.v += self.dv*dt
self.q = se3.integrate(self.model, self.q, v_mean*dt)
self.t += dt
return self.q, self.v
def reset(self):
self.first_iter = True
def simulate(self, u, dt=0.001, ndt=1):
''' Perform ndt steps, each lasting dt/ndt seconds '''
tau_c_avg = 0*self.tau_c
for i in range(ndt):
self.q, self.v = self.step(u, dt/ndt)
tau_c_avg += self.tau_c
self.tau_c = tau_c_avg / ndt
return self.q, self.v, self.f
If you prefer I can push it in the repo, together with a simple test and the additional code for visualization with gepetto-viewer. I don't know if all of this is necessary though if we wanna use it just as inspiration for the design.