Update ABA and RNEA

flferretti · Nov 13, 2023 · fbe588b · fbe588b
1 parent 408c9e1
commit fbe588b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 22 deletions.
diff --git a/src/jaxsim/physics/algos/aba.py b/src/jaxsim/physics/algos/aba.py
@@ -40,7 +40,7 @@ def aba(
     IM = jnp.array(
         [jnp.eye(6) * m for m in [*model._joint_motor_inertia.values()]] * model.NB
     )
-    K̅ᵥ = Γ.T @ jnp.array([*model._joint_motor_viscous_friction.values()]) @ Γ
+    K̅ᵥ = Γ.T * jnp.array([*model._joint_motor_viscous_friction.values()]) * Γ
     m_S = jnp.concatenate([S[:1], S[1:] * Γ[:, None, None]], axis=0)
 
     # Initialize buffers
@@ -136,7 +136,7 @@ def loop_body_pass1(carry: Pass1Carry, i: jtp.Int) -> Tuple[Pass1Carry, None]:
         pA_i = Cross.vx_star(v[i]) @ M[i] @ v[i] - i_Xf_W @ jnp.vstack(f_ext[i])
         pA = pA.at[i].set(pA_i)
 
-        pR_i = Cross.vx_star(m_v[i]) @ IM[i] @ m_v[i] + K̅ᵥ[i] * m_v[i]
+        pR_i = Cross.vx_star(m_v[i]) @ IM[i] @ m_v[i] - K̅ᵥ[i] * m_v[i]
         pR = pR.at[i].set(pR_i)
 
         return (i_X_λi, v, c, m_v, m_c, MA, pA, pR, i_X_0), None
@@ -174,7 +174,7 @@ def loop_body_pass2(carry: Pass2Carry, i: jtp.Int) -> Tuple[Pass2Carry, None]:
         u_i = tau[ii] - S[i].T @ pA[i] if tau.size != 0 else -S[i].T @ pA[i]
         u = u.at[i].set(u_i.squeeze())
 
-        has_motors = (Γ[i] != 1).any()
+        has_motors = Γ[i] != 1
 
         m_u_i = (
             tau[ii] / Γ[i] * has_motors - m_S[i].T @ pR[i]

diff --git a/src/jaxsim/physics/algos/rnea.py b/src/jaxsim/physics/algos/rnea.py
@@ -41,23 +41,11 @@ def rnea(
     S = model.motion_subspaces(q=q)
     i_X_λi = jnp.zeros_like(pre_X_λi)
 
-    Γ = (
-        jnp.array(list(model._joint_motor_gear_ratio.values()))
-        if hasattr(model, "_joint_motor_gear_ratio")
-        else jnp.ones(model.dofs)
-    )
-    IM = (
-        jnp.array(list(model._joint_motor_inertia.values()))
-        if hasattr(model, "_joint_motor_inertia")
-        else jnp.zeros(model.dofs)
-    )
-    K_v = (
-        jnp.array(list(model._joint_motor_viscous_friction.values()))
-        if hasattr(model, "_joint_motor_viscous_friction")
-        else jnp.zeros(model.dofs)
-    )
+    Γ = jnp.array([*model._joint_motor_gear_ratio.values()])
+    IM = jnp.array([*model._joint_motor_inertia.values()])
+    K_v = jnp.array([*model._joint_motor_viscous_friction.values()])
     K̅ᵥ = jnp.diag(Γ.T * jnp.diag(K_v) * Γ)
-    S_m = jnp.concatenate([S[:1], S[1:] * Γ[:, None, None]], axis=0)
+    m_S = jnp.concatenate([S[:1], S[1:] * Γ[:, None, None]], axis=0)
 
     i_X_0 = jnp.zeros_like(pre_X_λi)
     i_X_0 = i_X_0.at[0].set(jnp.eye(6))
@@ -121,7 +109,7 @@ def forward_pass(
         i_X_λi, v, v_m, a, a_m, i_X_0, f, f_m = carry
 
         vJ = S[i] * qd[ii]
-        vJ_m = S_m[i] * qd[ii]
+        vJ_m = m_S[i] * qd[ii]
 
         i_X_λi_i = i_X_pre[i] @ pre_X_λi[i]
         i_X_λi = i_X_λi.at[i].set(i_X_λi_i)
@@ -135,7 +123,7 @@ def forward_pass(
         a_i = i_X_λi[i] @ a[λ[i]] + S[i] * qdd[ii] + Cross.vx(v[i]) @ vJ
         a = a.at[i].set(a_i)
 
-        a_i_m = i_X_λi[i] @ a_m[λ[i]] + S_m[i] * qdd[ii] + Cross.vx(v_m[i]) @ vJ_m
+        a_i_m = i_X_λi[i] @ a_m[λ[i]] + m_S[i] * qdd[ii] + Cross.vx(v_m[i]) @ vJ_m
         a_m = a_m.at[i].set(a_i_m)
 
         i_X_0_i = i_X_λi[i] @ i_X_0[λ[i]]
@@ -171,7 +159,7 @@ def backward_pass(
         ii = i - 1
         tau, f, f_m = carry
 
-        value = S[i].T @ f[i] + S_m[i].T @ f_m[i]  # + K̅ᵥ[i] * qd[ii]
+        value = S[i].T @ f[i] + m_S[i].T @ f_m[i]  # + K̅ᵥ[i] * qd[ii]
         tau = tau.at[ii].set(value.squeeze())
 
         def update_f(ffm: Tuple[jtp.MatrixJax, jtp.MatrixJax]) -> jtp.MatrixJax: