fixed inverse dynamcs formulation for go2

Author: lorenzo96-cmd

mpx/config/config_go2_inv.py

@@ -9,6 +9,7 @@
 from mujoco.mjx._src.dataclasses import PyTreeNode
 
 from mpx.jax_ocp_solvers.jax_ocp_solvers import optimizers
+import mpx.utils.mpc_utils as mpc_utils
 
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
@@ -21,7 +22,8 @@
 N = 12
 mpc_frequency = 50.0
 solver_mode = "equality"
-equality_num_alpha = 5
+equality_num_alpha = 10
+regularization = jnp.array(1e-6)
 
 timer_t = jnp.array([0.5, 0.0, 0.0, 0.5])
 duty_factor = jnp.array([0.65])
@@ -30,6 +32,7 @@
 initial_height = 0.1
 robot_height = 0.27
 use_terrain_estimation = False
+clearance_speed = 0.2
 
 p0 = jnp.array([0.0, 0.0, robot_height])
 quat0 = jnp.array([1.0, 0.0, 0.0, 0.0])
@@ -65,16 +68,28 @@
 tau_slice = slice(nv, nv + n_joints)
 grf_slice = slice(nv + n_joints, m)
 
-Qp = jnp.diag(jnp.array([0.0, 0.0, 1.0e4]))
+Qp = jnp.diag(jnp.array([0.0, 0.0, 1.0e6]))
 Qrot = jnp.diag(jnp.array([1.0e3, 1.0e3, 0.0]))
-Qq = jnp.diag(jnp.ones(n_joints)) * 1.0e0
+Qq = jnp.diag(jnp.ones(n_joints)) * 1.0e2
 Qdp = jnp.diag(jnp.array([1.0, 1.0, 1.0])) * 5.0e3
 Qomega = jnp.diag(jnp.array([1.0, 1.0, 1.0])) * 1.0e2
-Qdq = jnp.diag(jnp.ones(n_joints)) * 1.0e0
+Qdq = jnp.diag(jnp.ones(n_joints)) * 1.0e-1
+Qcontact = jnp.diag(jnp.tile(jnp.array([1.0e5, 1.0e5, 1.0e5]), n_contact))
 Qacc = jnp.diag(jnp.ones(nv)) * 1.0e-1
-Qtau = jnp.diag(jnp.ones(n_joints)) * 1.0e-3
-Qgrf = jnp.diag(jnp.ones(3 * n_contact)) * 1.0e-5
-W = jax.scipy.linalg.block_diag(Qp, Qrot, Qq, Qdp, Qomega, Qdq, Qacc, Qtau, Qgrf)
+Qtau = jnp.diag(jnp.ones(n_joints)* 1.0e-1) 
+Qgrf = jnp.diag(jnp.ones(3 * n_contact)) * 1.0e-2
+W = {
+    "pos": Qp,
+    "rot": Qrot,
+    "q": Qq,
+    "vel": Qdp,
+    "omega": Qomega,
+    "dq": Qdq,
+    "contact": Qcontact,
+    "acc": Qacc,
+    "tau": Qtau,
+    "grf": Qgrf,
+}
 
 initial_state = jnp.concatenate([p0, quat0, q0, jnp.zeros(nv)])
 
@@ -120,13 +135,20 @@ def dynamics(x, u, t, parameter):
     qpos_next, qvel_next = _integrate_state(qpos, qvel, qacc)
     return jnp.concatenate([qpos_next, qvel_next])
 
+@jax.jit
+def _foot_positions(qpos, qvel):
+    data = mjx.make_data(_MJX_MODEL)
+    data = data.replace(qpos=qpos, qvel=qvel)
+    data = mjx.fwd_position(_MJX_MODEL, data)
+    return jnp.concatenate([data.geom_xpos[contact_id] for contact_id in _CONTACT_ID])
+
 def equality(x, u, t, parameter):
-    del t
+    
     qpos, qvel = _state_parts(x)
     qacc = u[qacc_slice]
     tau = u[tau_slice]
     grf = u[grf_slice]
-    data = mjx.make_data(_MODEL)
+    data = mjx.make_data(_MJX_MODEL)
     data = data.replace(qpos=qpos, qvel=qvel, qacc=qacc)
     data = mjx.fwd_position(_MJX_MODEL, data)
     data = mjx.fwd_velocity(_MJX_MODEL, data)
@@ -138,9 +160,11 @@ def equality(x, u, t, parameter):
         jac, _ = mjx.jac(_MJX_MODEL, data, data.geom_xpos[contact_id], body_id)
         jacobians.append(jac)
     contact_jacobian = jnp.concatenate(jacobians, axis=1)
+    contact = parameter[t, :4]
+    grf = jnp.concatenate([grf[:3]*contact[0], grf[3:6]*contact[1], grf[6:9]*contact[2], grf[9:12]*contact[3]])
     contact_wrench = contact_jacobian @ grf
     generalized_actuation = jnp.concatenate([jnp.zeros(6, dtype=u.dtype), tau])
-    return qfrc_inverse + contact_wrench - generalized_actuation
+    return qfrc_inverse - contact_wrench - generalized_actuation
 
 
 def cost(W, reference, x, u, t):
@@ -154,85 +178,43 @@ def cost(W, reference, x, u, t):
     acc = u[qacc_slice]
     tau = u[tau_slice]
     grf = u[grf_slice]
-
-    p_ref = reference[t, :3]
-    quat_ref = reference[t, 3:7]
-    q_ref = reference[t, 7 : 7 + n_joints]
-    dp_ref = reference[t, nq : nq + 3]
-    omega_ref = reference[t, nq + 3 : nq + 6]
-    dq_ref = reference[t, nq + 6 : nq + nv]
+    p_leg = _foot_positions(qpos, qvel)
+    p_ref = reference["p"][t]
+    quat_ref = reference["quat"][t]
+    q_ref = reference["q"][t]
+    dp_ref = reference["dp"][t]
+    omega_ref = reference["omega"][t]
+    p_leg_ref = reference["foot"][t]
+    grf_ref = reference["grf"][t]
 
     quat_err = math.quat_sub(quat, quat_ref)
+    
     stage_cost = (
-        (p - p_ref).T @ W[:3, :3] @ (p - p_ref)
-        + quat_err.T @ W[3:6, 3:6] @ quat_err
-        + (q - q_ref).T @ W[6 : 6 + n_joints, 6 : 6 + n_joints] @ (q - q_ref)
-        + (dp - dp_ref).T
-        @ W[6 + n_joints : 9 + n_joints, 6 + n_joints : 9 + n_joints]
-        @ (dp - dp_ref)
-        + (omega - omega_ref).T
-        @ W[9 + n_joints : 12 + n_joints, 9 + n_joints : 12 + n_joints]
-        @ (omega - omega_ref)
-        + (dq - dq_ref).T
-        @ W[12 + n_joints : 12 + 2 * n_joints, 12 + n_joints : 12 + 2 * n_joints]
-        @ (dq - dq_ref)
-        + acc.T @ W[nx_error : nx_error + nv, nx_error : nx_error + nv] @ acc
-        + (tau ).T
-        @ W[nx_error + nv : nx_error + nv + n_joints, nx_error + nv : nx_error + nv + n_joints]
-        @ (tau)
-        + (grf).T
-        @ W[nx_error + nv + n_joints :, nx_error + nv + n_joints :]
-        @ (grf)
+        (p - p_ref).T @ W["pos"] @ (p - p_ref)
+        + quat_err.T @ W["rot"] @ quat_err
+        + (q - q_ref).T @ W["q"] @ (q - q_ref)
+        + (dp - dp_ref).T @ W["vel"] @ (dp - dp_ref)
+        + (omega - omega_ref).T @ W["omega"] @ (omega - omega_ref)
+        + dq.T @ W["dq"] @ dq
+        + (p_leg-p_leg_ref).T @ W["contact"] @ (p_leg-p_leg_ref)
+        + acc.T @ W["acc"] @ acc
+        + tau.T @ W["tau"] @ tau
+        + (grf - grf_ref).T @ W["grf"] @ (grf - grf_ref)
     )
     terminal_cost = (
-        (p - p_ref).T @ W[:3, :3] @ (p - p_ref)
-        + quat_err.T @ W[3:6, 3:6] @ quat_err
-        + (q - q_ref).T @ W[6 : 6 + n_joints, 6 : 6 + n_joints] @ (q - q_ref)
-        + (dp - dp_ref).T
-        @ W[6 + n_joints : 9 + n_joints, 6 + n_joints : 9 + n_joints]
-        @ (dp - dp_ref)
-        + (omega - omega_ref).T
-        @ W[9 + n_joints : 12 + n_joints, 9 + n_joints : 12 + n_joints]
-        @ (omega - omega_ref)
-        + (dq - dq_ref).T
-        @ W[12 + n_joints : 12 + 2 * n_joints, 12 + n_joints : 12 + 2 * n_joints]
-        @ (dq - dq_ref)
+        (p - p_ref).T @ W["pos"] @ (p - p_ref)
+        + quat_err.T @ W["rot"] @ quat_err
+        + (q - q_ref).T @ W["q"] @ (q - q_ref)
+        + (dp - dp_ref).T @ W["vel"] @ (dp - dp_ref)
+        + (omega - omega_ref).T @ W["omega"] @ (omega - omega_ref)
+        + dq.T @ W["dq"] @ dq
     )
     return jnp.where(t == N, 0.5 * terminal_cost, 0.5 * stage_cost)
 
 
-def hessian_approx(W, reference, x, u, t):
-    del reference, x, u
-    state_diag = jnp.concatenate(
-        [
-            jnp.diag(W[:3, :3]),
-            jnp.ones(4) * 1.0e-3,
-            jnp.diag(W[6 : 6 + n_joints, 6 : 6 + n_joints]),
-            jnp.diag(W[6 + n_joints : 12 + 2 * n_joints, 6 + n_joints : 12 + 2 * n_joints]),
-        ]
-    )
-    control_diag = jnp.concatenate(
-        [
-            jnp.diag(W[nx_error : nx_error + nv, nx_error : nx_error + nv]),
-            jnp.diag(
-                W[
-                    nx_error + nv : nx_error + nv + n_joints,
-                    nx_error + nv : nx_error + nv + n_joints,
-                ]
-            ),
-            jnp.diag(W[nx_error + nv + n_joints :, nx_error + nv + n_joints :]),
-        ]
-    )
-    control_diag = jnp.where(t == N, jnp.ones(m) * 1.0e-6, control_diag)
-
-    Q = jnp.diag(state_diag)
-    R = jnp.diag(control_diag)
-    M = jnp.zeros((n, m), dtype=W.dtype)
-    return Q, R, M
-
-
 class InverseDynamicsMPCData(PyTreeNode):
     dt: float
+    time: jnp.ndarray
     duty_factor: jnp.ndarray
     step_freq: jnp.ndarray
     step_height: jnp.ndarray
@@ -243,6 +225,7 @@ class InverseDynamicsMPCData(PyTreeNode):
     V0: jnp.ndarray
     Veq0: jnp.ndarray
     W: jnp.ndarray
+    regularization: jnp.ndarray
 
 
 @partial(jax.jit, static_argnums=(0, 1))
@@ -289,28 +272,43 @@ def __init__(self, config, limited_memory=False):
         self.model = mujoco.MjModel.from_xml_path(config.model_path)
         self.model.opt.timestep = config.dt
         self.mjx_model = mjx.put_model(self.model)
-
+        data = mujoco.MjData(self.model)
+        mujoco.mj_fwdPosition(self.model, data)
         self.initial_state = jnp.asarray(config.initial_state)
         self.initial_X0 = jnp.tile(self.initial_state, (config.N + 1, 1))
         self.initial_U0 = jnp.tile(config.u_ref, (config.N, 1))
         self.initial_V0 = jnp.zeros((config.N + 1, config.n))
         self.initial_Veq0 = jnp.zeros((config.N, config.equality_dim))
+        self.initial_liftoff = config.p_legs0
 
         self.dynamics = config.dynamics
         solver = partial(
             optimizers.mpc_equality,
             config.cost,
             self.dynamics,
-            config.hessian_approx,
+            None,
             limited_memory,
             equality=config.equality,
             num_alpha=config.equality_num_alpha,
         )
 
-        def solve(reference, parameter, W, x0, X0, U0, V0, Veq0):
-            return solver(reference, parameter, W, x0, X0, U0, V0, Veq_in=Veq0)
+        def solve(reference, parameter, W, x0, X0, U0, V0, Veq0, regularization):
+            return solver(
+                reference,
+                parameter,
+                W,
+                x0,
+                X0,
+                U0,
+                V0,
+                Veq_in=Veq0,
+                regularization=regularization,
+            )
 
         self._solve = jax.jit(solve)
+        robot_mass = data.qM[0]
+        self._ref_gen = partial(config.reference_generator, mass=robot_mass)
+        self._timer_run = jax.jit(mpc_utils.timer_run)
         self._update_warm_start = partial(
             _update_warm_start,
             config.N,
@@ -321,37 +319,50 @@ def solve(reference, parameter, W, x0, X0, U0, V0, Veq0):
     def make_data(self):
         return InverseDynamicsMPCData(
             dt=dt,
+            time=jnp.asarray(0.0, dtype=jnp.float32),
             duty_factor=duty_factor,
             step_freq=step_freq,
             step_height=step_height,
             contact_time=timer_t,
-            liftoff=jnp.zeros(3 * n_contact),
+            liftoff=self.initial_liftoff,
             X0=self.initial_X0,
             U0=self.initial_U0,
             V0=self.initial_V0,
             Veq0=self.initial_Veq0,
             W=W,
+            regularization=regularization,
         )
 
-    def _reference(self, x, command):
-        p_ref = x[:3] + jnp.array([command[0], command[1], 0.0]) * (N * dt)
-        p_ref = p_ref.at[2].set(command[6])
-        qpos_ref = jnp.concatenate([p_ref, x[3:7], q0])
-        dp_ref = jnp.array([command[0], command[1], 0.0])
-        omega_ref = jnp.array([0.0, 0.0, command[5]])
-        dq_ref = jnp.zeros(n_joints)
-        x_ref = jnp.concatenate([qpos_ref, dp_ref, omega_ref, dq_ref])
-        return jnp.tile(x_ref, (N + 1, 1))
-
     def control_output(self, x0, X, U, reference, parameter):
         del x0, X, reference, parameter
         return jnp.clip(U[0, tau_slice], min_torque, max_torque)
 
     def run(self, data, x, command, contact=None):
-        del contact
-        reference = self._reference(x, command)
-        parameter = jnp.zeros((N + 1, 1), dtype=x.dtype)
-        X, U, V, Veq = self._solve(
+        if contact is None:
+            contact = jnp.zeros(n_contact, dtype=x.dtype)
+        current_time = data.time + jnp.asarray(1 / self.mpc_frequency, dtype=data.time.dtype)
+        _, contact_time = self._timer_run(
+            data.duty_factor,
+            data.step_freq,
+            data.contact_time,
+            1 / self.mpc_frequency,
+        )
+        qpos, qvel = _state_parts(x)
+        foot = _foot_positions(qpos, qvel)
+        reference_state = jnp.concatenate([x, foot, jnp.zeros(3 * n_contact, dtype=x.dtype)])
+        reference, parameter, liftoff = self._ref_gen(
+            duty_factor=data.duty_factor,
+            step_freq=data.step_freq,
+            step_height=data.step_height,
+            t_timer=data.contact_time,
+            x=reference_state,
+            foot=foot,
+            input=command,
+            liftoff=data.liftoff,
+            contact=contact,
+            current_time=current_time,
+        )
+        X, U, V, Veq, regularization, alpha_best, any_accepted = self._solve(
             reference,
             parameter,
             data.W,
@@ -360,6 +371,7 @@ def run(self, data, x, command, contact=None):
             data.U0,
             data.V0,
             data.Veq0,
+            data.regularization,
         )
         valid_solution = jnp.logical_not(jnp.isnan(U[0, 0]))
         tau = jax.lax.cond(
@@ -377,23 +389,54 @@ def run(self, data, x, command, contact=None):
             V,
             Veq,
         )
-        return data.replace(X0=X0, U0=U0, V0=V0, Veq0=Veq0), tau, q, dq
+        return (
+            data.replace(
+                time=current_time,
+                X0=X0,
+                U0=U0,
+                V0=V0,
+                Veq0=Veq0,
+                contact_time=contact_time,
+                liftoff=liftoff,
+                regularization=regularization,
+            ),
+            tau,
+            q,
+            dq,
+            alpha_best,
+            any_accepted,
+            reference,
+        )
 
     def reset(self, data, qpos, qvel, foot=None):
-        del foot
         x = (
             self.initial_state.at[self.qpos_slice].set(jnp.ravel(qpos))
             .at[self.qvel_slice].set(jnp.ravel(qvel))
         )
+        liftoff = _foot_positions(jnp.ravel(qpos), jnp.ravel(qvel)) if foot is None else jnp.ravel(foot)
         return data.replace(
+            time=jnp.asarray(0.0, dtype=jnp.float32),
             X0=jnp.tile(x, (N + 1, 1)),
             U0=self.initial_U0,
             V0=self.initial_V0,
             Veq0=self.initial_Veq0,
-            liftoff=jnp.zeros(3 * n_contact),
+            liftoff=liftoff,
             contact_time=timer_t,
         )
 
 
 def state_to_qpos(x):
     return x[:nq]
+
+
+reference_generator = partial(
+    mpc_utils.reference_generator,
+    use_terrain_estimation,
+    N,
+    dt,
+    n_joints,
+    n_contact,
+    foot0=p_legs0,
+    q0=q0,
+    clearence_speed=clearance_speed,
+)