[Feature] Joint friction / frictionloss

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -410,6 +410,7 @@ def _load_scene(self, morph, surface):
             j_info["dofs_limit"] = np.tile([-np.inf, np.inf], (6, 1))
             j_info["dofs_stiffness"] = np.zeros(6)
             j_info["dofs_invweight"] = np.zeros(6)
+            j_info["dofs_frictionloss"] = np.zeros(6)
             j_info["dofs_damping"] = np.zeros(6)
             if isinstance(morph, gs.morphs.Drone):
                 mass_tot = sum(l_info["inertial_mass"] for l_info in l_infos)
@@ -697,6 +698,7 @@ def _add_by_info(self, l_info, j_infos, g_infos, morph, surface):
                 dofs_motion_vel=dofs_motion_vel,
                 dofs_limit=j_info.get("dofs_limit", np.tile([[-np.inf, np.inf]], [n_dofs, 1])),
                 dofs_invweight=j_info.get("dofs_invweight", np.zeros(n_dofs)),
+                dofs_frictionloss=j_info.get("dofs_frictionloss", np.zeros(n_dofs)),
                 dofs_stiffness=j_info.get("dofs_stiffness", np.zeros(n_dofs)),
                 dofs_damping=j_info.get("dofs_damping", np.zeros(n_dofs)),
                 dofs_armature=j_info.get("dofs_armature", np.zeros(n_dofs)),

genesis/engine/entities/rigid_entity/rigid_joint.py

@@ -32,6 +32,7 @@ def __init__(
         dofs_motion_vel,
         dofs_limit,
         dofs_invweight,
+        dofs_frictionloss,
         dofs_stiffness,
         dofs_damping,
         dofs_armature,
@@ -60,6 +61,7 @@ def __init__(
         self._dofs_motion_vel = dofs_motion_vel
         self._dofs_limit = dofs_limit
         self._dofs_invweight = dofs_invweight
+        self._dofs_frictionloss = dofs_frictionloss
         self._dofs_stiffness = dofs_stiffness
         self._dofs_damping = dofs_damping
         self._dofs_armature = dofs_armature
@@ -390,6 +392,13 @@ def dofs_invweight(self):
         """
         return self._dofs_invweight
 
+    @property
+    def dofs_frictionloss(self):
+        """
+        Returns the frictionloss of the dofs of the joint.
+        """
+        return self._dofs_frictionloss
+
     @property
     def dofs_stiffness(self):
         """

genesis/engine/solvers/rigid/constraint_solver_decomp.py

@@ -117,6 +117,7 @@ def __init__(self, rigid_solver: "RigidSolver"):
         self.jac_relevant_dofs = cs.jac_relevant_dofs
         self.n_constraints = cs.n_constraints
         self.n_constraints_equality = cs.n_constraints_equality
+        self.n_constraints_frictionloss = cs.n_constraints_frictionloss
         self.improved = cs.improved
         self.Jaref = cs.Jaref
         self.Ma = cs.Ma
@@ -155,52 +156,6 @@ def __init__(self, rigid_solver: "RigidSolver"):
 
         self.reset()
 
-        #
-
-        # self.constraint_state.ti_n_equalities = self.ti_n_equalities
-        # self.constraint_state.jac = self.jac
-        # self.constraint_state.diag = self.diag
-        # self.constraint_state.aref = self.aref
-        # self.constraint_state.jac_n_relevant_dofs = self.jac_n_relevant_dofs
-        # self.constraint_state.jac_relevant_dofs = self.jac_relevant_dofs
-        # self.constraint_state.n_constraints = self.n_constraints
-        # self.constraint_state.n_constraints_equality = self.n_constraints_equality
-        # self.constraint_state.improved = self.improved
-        # self.constraint_state.Jaref = self.Jaref
-        # self.constraint_state.Ma = self.Ma
-        # self.constraint_state.Ma_ws = self.Ma_ws
-        # self.constraint_state.grad = self.grad
-        # self.constraint_state.Mgrad = self.Mgrad
-        # self.constraint_state.search = self.search
-        # self.constraint_state.efc_D = self.efc_D
-        # self.constraint_state.efc_force = self.efc_force
-        # self.constraint_state.active = self.active
-        # self.constraint_state.prev_active = self.prev_active
-        # self.constraint_state.qfrc_constraint = self.qfrc_constraint
-        # self.constraint_state.qacc = self.qacc
-        # self.constraint_state.qacc_ws = self.qacc_ws
-        # self.constraint_state.qacc_prev = self.qacc_prev
-        # self.constraint_state.cost_ws = self.cost_ws
-        # self.constraint_state.gauss = self.gauss
-        # self.constraint_state.cost = self.cost
-        # self.constraint_state.prev_cost = self.prev_cost
-        # self.constraint_state.gtol = self.gtol
-        # self.constraint_state.mv = self.mv
-        # self.constraint_state.jv = self.jv
-        # self.constraint_state.quad_gauss = self.quad_gauss
-        # self.constraint_state.quad = self.quad
-        # self.constraint_state.candidates = self.candidates
-        # self.constraint_state.ls_its = self.ls_its
-        # self.constraint_state.ls_result = self.ls_result
-        # if self._solver_type == gs.constraint_solver.CG:
-        #     self.constraint_state.cg_prev_grad = self.cg_prev_grad
-        #     self.constraint_state.cg_prev_Mgrad = self.cg_prev_Mgrad
-        #     self.constraint_state.cg_beta = self.cg_beta
-        #     self.constraint_state.cg_pg_dot_pMg = self.cg_pg_dot_pMg
-        # if self._solver_type == gs.constraint_solver.Newton:
-        #     self.constraint_state.nt_H = self.nt_H
-        #     self.constraint_state.nt_vec = self.nt_vec
-
     def clear(self, envs_idx: npt.NDArray[np.int32] | None = None):
         if envs_idx is None:
             envs_idx = self._solver._scene._envs_idx
@@ -229,6 +184,16 @@ def handle_constraints(self):
             static_rigid_sim_config=self._solver._static_rigid_sim_config,
         )
 
+        add_frictionloss_constraints(
+            links_info=self._solver.links_info,
+            joints_info=self._solver.joints_info,
+            dofs_info=self._solver.dofs_info,
+            dofs_state=self._solver.dofs_state,
+            rigid_global_info=self._solver._rigid_global_info,
+            constraint_state=self.constraint_state,
+            static_rigid_sim_config=self._solver._static_rigid_sim_config,
+        )
+
         if self._solver._enable_collision:
             add_collision_constraints(
                 links_info=self._solver.links_info,
@@ -300,6 +265,7 @@ def constraint_solver_kernel_clear(
         i_b = envs_idx[i_b_]
         constraint_state.n_constraints[i_b] = 0
         constraint_state.n_constraints_equality[i_b] = 0
+        constraint_state.n_constraints_frictionloss[i_b] = 0
 
 
 @ti.kernel
@@ -804,7 +770,7 @@ def func_equality_weld(
     for i_ab in range(2):
         sign = gs.ti_float(1.0) if i_ab == 0 else gs.ti_float(-1.0)
         link = link1_idx if i_ab == 0 else link2_idx
-        # For rotation, we use the body’s orientation (here we use its quaternion)
+        # For rotation, we use the body's orientation (here we use its quaternion)
         # and a suitable reference frame. (You may need a more detailed implementation.)
         while link > -1:
             link_maybe_batch = [link, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else link
@@ -906,6 +872,50 @@ def add_joint_limit_constraints(
                             constraint_state.jac_relevant_dofs[n_con, 0, i_b] = i_d
 
 
+@ti.kernel
+def add_frictionloss_constraints(
+    links_info: array_class.LinksInfo,
+    joints_info: array_class.JointsInfo,
+    dofs_info: array_class.DofsInfo,
+    dofs_state: array_class.DofsState,
+    rigid_global_info: array_class.RigidGlobalInfo,
+    constraint_state: array_class.ConstraintState,
+    static_rigid_sim_config: ti.template(),
+):
+    _B = constraint_state.jac.shape[2]
+    n_links = links_info.root_idx.shape[0]
+    n_dofs = dofs_state.ctrl_mode.shape[0]
+
+    # TODO: sparse mode
+    ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.PARTIAL)
+    for i_b in range(_B):
+        for i_l in range(n_links):
+            I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l
+
+            for i_j in range(links_info.joint_start[I_l], links_info.joint_end[I_l]):
+                I_j = [i_j, i_b] if ti.static(static_rigid_sim_config.batch_joints_info) else i_j
+
+                for i_d in range(joints_info.dof_start[I_j], joints_info.dof_end[I_j]):
+                    I_d = [i_d, i_b] if ti.static(static_rigid_sim_config.batch_dofs_info) else i_d
+
+                    if dofs_info.frictionloss[I_d] > 0:
+                        jac = 1.0
+                        jac_qvel = jac * dofs_state.vel[i_d, i_b]
+                        imp, aref = gu.imp_aref(joints_info.sol_params[I_j], 0.0, jac_qvel, 0.0)
+                        diag = ti.max(dofs_info.invweight[I_d] * (1 - imp) / imp, gs.EPS)
+
+                        n_con = constraint_state.n_constraints[i_b]
+                        ti.atomic_add(constraint_state.n_constraints_frictionloss[i_b], 1)
+                        constraint_state.n_constraints[i_b] = n_con + 1
+                        constraint_state.diag[n_con, i_b] = diag
+                        constraint_state.aref[n_con, i_b] = aref
+                        constraint_state.efc_D[n_con, i_b] = 1 / diag
+                        constraint_state.efc_frictionloss[n_con, i_b] = dofs_info.frictionloss[I_d]
+                        for i_d2 in range(n_dofs):
+                            constraint_state.jac[n_con, i_d2, i_b] = gs.ti_float(0.0)
+                        constraint_state.jac[n_con, i_d, i_b] = jac
+
+
 @ti.func
 def func_nt_hessian_incremental(
     i_b,
@@ -1146,7 +1156,9 @@ def func_update_contact_force(
 
     ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
     for i_b in range(_B):
-        const_start = constraint_state.n_constraints_equality[i_b]
+        const_start = constraint_state.n_constraints_equality[i_b] + constraint_state.n_constraints_frictionloss[i_b]
+
+        # contact constraints should be after equality and frictionloss constraints and before joint limit constraints
         for i_c in range(collider_state.n_contacts[i_b]):
             contact_data_normal = collider_state.contact_data.normal[i_c, i_b]
             contact_data_friction = collider_state.contact_data.friction[i_c, i_b]
@@ -1290,17 +1302,38 @@ def func_ls_point_fn(
     constraint_state: array_class.ConstraintState,
 ):
     tmp_quad_total0, tmp_quad_total1, tmp_quad_total2 = gs.ti_float(0.0), gs.ti_float(0.0), gs.ti_float(0.0)
-    for _i0 in range(1):
-        tmp_quad_total0 = constraint_state.quad_gauss[_i0 + 0, i_b]
-        tmp_quad_total1 = constraint_state.quad_gauss[_i0 + 1, i_b]
-        tmp_quad_total2 = constraint_state.quad_gauss[_i0 + 2, i_b]
-        for i_c in range(constraint_state.n_constraints[i_b]):
-            active = 1
-            if i_c >= constraint_state.n_constraints_equality[i_b]:
-                active = constraint_state.Jaref[i_c, i_b] + alpha * constraint_state.jv[i_c, i_b] < 0
-            tmp_quad_total0 += constraint_state.quad[i_c, _i0 + 0, i_b] * active
-            tmp_quad_total1 += constraint_state.quad[i_c, _i0 + 1, i_b] * active
-            tmp_quad_total2 += constraint_state.quad[i_c, _i0 + 2, i_b] * active
+    tmp_quad_total0 = constraint_state.quad_gauss[0, i_b]
+    tmp_quad_total1 = constraint_state.quad_gauss[1, i_b]
+    tmp_quad_total2 = constraint_state.quad_gauss[2, i_b]
+    for i_c in range(constraint_state.n_constraints[i_b]):
+        x = constraint_state.Jaref[i_c, i_b] + alpha * constraint_state.jv[i_c, i_b]
+        active = 1
+        ne = constraint_state.n_constraints_equality[i_b]
+        nef = ne + constraint_state.n_constraints_frictionloss[i_b]
+        qf_0 = constraint_state.quad[i_c, 0, i_b]
+        qf_1 = constraint_state.quad[i_c, 1, i_b]
+        qf_2 = constraint_state.quad[i_c, 2, i_b]
+        if i_c >= ne and i_c < nef:
+            f = constraint_state.efc_frictionloss[i_c, i_b]
+            r = 1.0 / ti.max(constraint_state.efc_D[i_c, i_b], gs.EPS)
+            rf = r * f
+            linear_neg = x <= -rf
+            linear_pos = x >= rf
+
+            if linear_neg or linear_pos:
+                qf_0 = linear_neg * f * (-0.5 * rf - constraint_state.Jaref[i_c, i_b]) + linear_pos * f * (
+                    -0.5 * rf + constraint_state.Jaref[i_c, i_b]
+                )
+                qf_1 = linear_neg * (-f * constraint_state.jv[i_c, i_b]) + linear_pos * (
+                    f * constraint_state.jv[i_c, i_b]
+                )
+                qf_2 = 0.0
+        elif i_c >= nef:
+            active = x < 0
+
+        tmp_quad_total0 += qf_0 * active
+        tmp_quad_total1 += qf_1 * active
+        tmp_quad_total2 += qf_2 * active
 
     cost = alpha * alpha * tmp_quad_total2 + alpha * tmp_quad_total1 + tmp_quad_total0
 
@@ -1550,8 +1583,6 @@ def update_bracket(
                 constraint_state.candidates[4 * i + 3, i_b],
             )
             flag = 2
-        else:
-            pass
 
     p_next_alpha, p_next_cost, p_next_deriv_0, p_next_deriv_1 = p_alpha, p_cost, p_deriv_0, p_deriv_1
 
@@ -1681,9 +1712,24 @@ def func_update_constraint(
         if ti.static(static_rigid_sim_config.solver_type == gs.constraint_solver.Newton):
             constraint_state.prev_active[i_c, i_b] = constraint_state.active[i_c, i_b]
         constraint_state.active[i_c, i_b] = 1
-        if i_c >= constraint_state.n_constraints_equality[i_b]:
+        floss_force = gs.ti_float(0.0)
+        ne = constraint_state.n_constraints_equality[i_b]
+        nef = ne + constraint_state.n_constraints_frictionloss[i_b]
+        if i_c >= ne and i_c < nef:
+            f = constraint_state.efc_frictionloss[i_c, i_b]
+            r = 1.0 / ti.max(constraint_state.efc_D[i_c, i_b], gs.EPS)
+            rf = r * f
+            linear_neg = constraint_state.Jaref[i_c, i_b] <= -rf
+            linear_pos = constraint_state.Jaref[i_c, i_b] >= rf
+            constraint_state.active[i_c, i_b] = (~linear_neg) & (~linear_pos)
+            floss_force = linear_neg * f + linear_pos * -f
+            floss_cost_local = linear_neg * f * (-0.5 * rf - constraint_state.Jaref[i_c, i_b])
+            floss_cost_local += linear_pos * f * (-0.5 * rf + constraint_state.Jaref[i_c, i_b])
+            cost[i_b] = cost[i_b] + floss_cost_local
+        elif i_c >= nef:
             constraint_state.active[i_c, i_b] = constraint_state.Jaref[i_c, i_b] < 0
-        constraint_state.efc_force[i_c, i_b] = (
+
+        constraint_state.efc_force[i_c, i_b] = floss_force + (
             -constraint_state.efc_D[i_c, i_b] * constraint_state.Jaref[i_c, i_b] * constraint_state.active[i_c, i_b]
         )
 

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -562,6 +562,7 @@ def _init_dof_fields(self):
                 dofs_invweight=np.concatenate([joint.dofs_invweight for joint in joints], dtype=gs.np_float),
                 dofs_stiffness=np.concatenate([joint.dofs_stiffness for joint in joints], dtype=gs.np_float),
                 dofs_damping=np.concatenate([joint.dofs_damping for joint in joints], dtype=gs.np_float),
+                dofs_frictionloss=np.concatenate([joint.dofs_frictionloss for joint in joints], dtype=gs.np_float),
                 dofs_armature=np.concatenate([joint.dofs_armature for joint in joints], dtype=gs.np_float),
                 dofs_kp=np.concatenate([joint.dofs_kp for joint in joints], dtype=gs.np_float),
                 dofs_kv=np.concatenate([joint.dofs_kv for joint in joints], dtype=gs.np_float),
@@ -2510,6 +2511,7 @@ def kernel_init_dof_fields(
     dofs_invweight: ti.types.ndarray(),
     dofs_stiffness: ti.types.ndarray(),
     dofs_damping: ti.types.ndarray(),
+    dofs_frictionloss: ti.types.ndarray(),
     dofs_armature: ti.types.ndarray(),
     dofs_kp: ti.types.ndarray(),
     dofs_kv: ti.types.ndarray(),
@@ -2538,6 +2540,7 @@ def kernel_init_dof_fields(
         dofs_info.invweight[I] = dofs_invweight[i]
         dofs_info.stiffness[I] = dofs_stiffness[i]
         dofs_info.damping[I] = dofs_damping[i]
+        dofs_info.frictionloss[I] = dofs_frictionloss[i]
         dofs_info.kp[I] = dofs_kp[i]
         dofs_info.kv[I] = dofs_kv[i]
 

genesis/utils/array_class.py

@@ -113,6 +113,7 @@ class StructConstraintState:
     jac_relevant_dofs: V_ANNOTATION
     jac_n_relevant_dofs: V_ANNOTATION
     n_constraints_equality: V_ANNOTATION
+    n_constraints_frictionloss: V_ANNOTATION
     improved: V_ANNOTATION
     Jaref: V_ANNOTATION
     Ma: V_ANNOTATION
@@ -121,6 +122,7 @@ class StructConstraintState:
     Mgrad: V_ANNOTATION
     search: V_ANNOTATION
     efc_D: V_ANNOTATION
+    efc_frictionloss: V_ANNOTATION
     efc_force: V_ANNOTATION
     active: V_ANNOTATION
     prev_active: V_ANNOTATION
@@ -171,6 +173,7 @@ def get_constraint_state(constraint_solver, solver):
         "jac_relevant_dofs": V(gs.ti_int, shape=solver._batch_shape((len_constraints_, solver.n_dofs_))),
         "jac_n_relevant_dofs": V(gs.ti_int, shape=solver._batch_shape(len_constraints_)),
         "n_constraints_equality": V(gs.ti_int, shape=solver._batch_shape()),
+        "n_constraints_frictionloss": V(gs.ti_int, shape=solver._batch_shape()),
         "improved": V(gs.ti_int, shape=solver._batch_shape()),
         "Jaref": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
         "Ma": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
@@ -179,6 +182,7 @@ def get_constraint_state(constraint_solver, solver):
         "Mgrad": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
         "search": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
         "efc_D": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
+        "efc_frictionloss": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
         "efc_force": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
         "active": V(dtype=gs.ti_int, shape=solver._batch_shape(len_constraints_)),
         "prev_active": V(dtype=gs.ti_int, shape=solver._batch_shape(len_constraints_)),
@@ -1067,6 +1071,7 @@ class StructDofsInfo:
     invweight: V_ANNOTATION
     armature: V_ANNOTATION
     damping: V_ANNOTATION
+    frictionloss: V_ANNOTATION
     motion_ang: V_ANNOTATION
     motion_vel: V_ANNOTATION
     limit: V_ANNOTATION
@@ -1083,6 +1088,7 @@ def get_dofs_info(solver):
         "invweight": V(dtype=gs.ti_float, shape=shape),
         "armature": V(dtype=gs.ti_float, shape=shape),
         "damping": V(dtype=gs.ti_float, shape=shape),
+        "frictionloss": V(dtype=gs.ti_float, shape=shape),
         "motion_ang": V(dtype=gs.ti_vec3, shape=shape),
         "motion_vel": V(dtype=gs.ti_vec3, shape=shape),
         "limit": V(dtype=gs.ti_vec2, shape=shape),

genesis/utils/mjcf.py

@@ -234,13 +234,12 @@ def parse_link(mj, i_l, scale):
         j_info["dofs_damping"] = mj.dof_damping[mj_dof_offset : (mj_dof_offset + n_dofs)]
         j_info["dofs_invweight"] = mj.dof_invweight0[mj_dof_offset : (mj_dof_offset + n_dofs)]
         j_info["dofs_armature"] = mj.dof_armature[mj_dof_offset : (mj_dof_offset + n_dofs)]
+        j_info["dofs_frictionloss"] = mj.dof_frictionloss[mj_dof_offset : (mj_dof_offset + n_dofs)]
         if mj.njnt > 0:
             mj_jnt_offset = i_j if i_j != -1 else 0
             j_info["sol_params"] = np.concatenate((mj.jnt_solref[mj_jnt_offset], mj.jnt_solimp[mj_jnt_offset]))
         else:
             j_info["sol_params"] = gu.default_solver_params()  # Placeholder. It will not be used anyway.
-        if (mj.dof_frictionloss[mj_dof_offset : (mj_dof_offset + n_dofs)] > 0.0).any():
-            gs.logger.warning("(MJCF) Friction loss at DoF-level not supported.")
 
         # Parsing joint parameters that are type-specific
         mj_stiffness = mj.jnt_stiffness[i_j] if i_j != -1 else 0.0
@@ -285,9 +284,6 @@ def parse_link(mj, i_l, scale):
 
                 j_info["init_qpos"] *= scale
 
-        if (mj.dof_frictionloss[mj_dof_offset : (mj_dof_offset + n_dofs)] > 0.0).any():
-            gs.logger.warning("(MJCF) Joint Coulomb friction not supported.")
-
         # Parsing actuator parameters
         j_info["dofs_kp"] = np.zeros((n_dofs,), dtype=gs.np_float)
         j_info["dofs_kv"] = np.zeros((n_dofs,), dtype=gs.np_float)

genesis/utils/urdf.py

@@ -262,6 +262,7 @@ def parse_urdf(morph, surface):
             gs.raise_exception(f"Unsupported URDF joint type: {joint.joint_type}")
 
         j_info["dofs_invweight"] = np.zeros(j_info["n_dofs"])
+        j_info["dofs_frictionloss"] = np.zeros(j_info["n_dofs"])
         j_info["sol_params"] = gu.default_solver_params()
         j_info["dofs_kp"] = gu.default_dofs_kp(j_info["n_dofs"])
         j_info["dofs_kv"] = gu.default_dofs_kv(j_info["n_dofs"])