How to calculate the matrices A and B of a system model for a robot in mujoco/python

Question

I want to calcultae the matrices A and B (xk+1 = Axk + Buk) from a robot model(kukaiiwa14) while simulation using mujoco_py. I need A and B matrices to compute optimal control on my robot. If someone has an idea how to solve this problem I'll be so grateful.

This is my code:

odel=load_model_from_path("/home/boudhina/.mujoco/mujoco210/kuka/kuka.xml")
sim=MjSim(model)
viewer=MjViewer(sim)

sim.step(1)
t=0
while True:
 ##simulation:
 viewer.render()
 t+=1
 sim.step(1)

 ##Variables:
 nq=sim.model.nq
 nv=sim.model.nv
 #sim.data.qfrc_applied[2]=5
 time=sim.data.time
 qpos=sim.data.qpos
 qvel=sim.data.qvel
 C=sim.data.qfrc_bias
 M=sim.data.qM
 Jc=sim.data.efc_J
 print('constraint Jacobian:',Jc)
 print('Size of constraint Jacobian:',np.shape(Jc))
 def Jacobian(n):
  J = np.zeros((6, nq))
  jacp = np.zeros(nv*3)
  jacr = np.zeros(nv*3)
  mjp.cymj._mj_jac(sim.model, sim.data, jacp, jacr, sim.data.get_body_xpos('iiwa_link_4'), 
sim.model.body_name2id('iiwa_link_4'))
  J[0:3,:] = jacr.reshape(3, nv)
  J[3:6,:] = jacp.reshape(3, nv)
  return J

 J=Jacobian(nq)
 print('Jacobian:',J)
 print('Size of Jacobian:',np.shape(J))
 wdot=np.dot(J,qvel)
 print('wdot:',wdot)
 '''wdot=task space position'''
 #sim.data.qfrc_inverse
 ''' _mj_fullM Convert sparse inertia matrix M into full (i.e. dense) matrix.'''
 def mass_matrix(nv):
  mass_matrix=np.ndarray(shape=(nv ** 2,), dtype=np.float64, order="C")
  mjp.cymj._mj_fullM(sim.model,mass_matrix,M)
  mass_matrix=np.reshape(mass_matrix,(nv,nv))
  return mass_matrix
 ##computed-torque-control:
 t=sim.data.time
 #kp=700
 #kv=400

 def coef_ubertragungsfkt(T1,T2):
  kp=1/(T1*T2)
  kv=(T1+T2)/(T1*T2)
  return kp,kv
 kp,kv=coef_ubertragungsfkt(0.02,0.5)



def coef(T,init_pos,end_pos):
  a_0=init_pos
  a_1=0
  a_3=-(2/(T*T*T))*(end_pos-a_0)
  a_2=-3/2*a_3*T
  return a_0,a_1,a_2,a_3


##trajectory_planning
 a_0,a_1,a_2,a_3=coef(4,1,2)
 #b_0,b_1,b_2,b_3=coef(4,2,3)
 for i in range(1,5,1):
  f=a_0+a_1*i+a_2*i**2+a_3*i**3
  g=a_1+2*a_2*i+3*a_3*i**2
  M1=mass_matrix(nv)
  qsoll=np.array([1,1.5,0,f,0,0,0])
  qvelsoll= np.array([0,0,0,g,0,0,0])
  qacc=kv*(qvelsoll-qvel)+kp*(qsoll-qpos)
  tau=np.dot(M1,qacc)+C

  sim.data.ctrl[1]=tau[0]
  sim.data.ctrl[4]=tau[1]
  sim.data.ctrl[7]=tau[2]
  sim.data.ctrl[10]=tau[3]
  sim.data.ctrl[13]=tau[4]
  sim.data.ctrl[17]=tau[5]
  sim.data.ctrl[20]=tau[6]
 if  t>4:
    break