Skip to content

[Feature] Joint friction / frictionloss #1479

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
duburcqa merged 14 commits into from
Aug 5, 2025
Merged

[Feature] Joint friction / frictionloss #1479

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
5a96162
frictionloss
YilingQiao Jul 24, 2025
7749d65
merge
YilingQiao Jul 28, 2025
c09e153
Merge branch 'main' into yiling/250728_frictionloss
YilingQiao Jul 30, 2025
4a20d83
wip frictionloss
YilingQiao Jul 30, 2025
a8c9343
add tests
YilingQiao Jul 31, 2025
10c0ed5
comments
YilingQiao Jul 31, 2025
63a6b97
clearn
YilingQiao Jul 31, 2025
196a718
Merge branch 'main' into yiling/250728_frictionloss
YilingQiao Jul 31, 2025
c3e34da
format
YilingQiao Jul 31, 2025
27f89a1
address comments
YilingQiao Jul 31, 2025
791d822
minor
YilingQiao Jul 31, 2025
ec6d455
Fix friction constraint gradient computation.
duburcqa Aug 5, 2025
246c8c6
Update unit test.
duburcqa Aug 5, 2025
95c000b
Merge branch 'main' into yiling/250728_frictionloss
duburcqa Aug 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions genesis/assets/xml/frictionloss_pendulum.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
<mujoco>

<default>
<joint damping=".01" frictionloss=".1"/>
</default>

<worldbody>
<light pos="0 0 1"/>
<geom type="plane" size="1 1 .01" pos="0 0 -1"/>
<body name="pendulum" pos="0.15 0 0">
<joint name="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
<body name="link1" pos="0.1 0 0">
<joint name="slide" type="slide" axis="1 0 0" stiffness="200"/>
<geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
</body>
</body>
</worldbody>

</mujoco>
2 changes: 2 additions & 0 deletions genesis/engine/entities/rigid_entity/rigid_entity.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -409,6 +409,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)
Expand Down Expand Up @@ -696,6 +697,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)),
Expand Down
9 changes: 9 additions & 0 deletions genesis/engine/entities/rigid_entity/rigid_joint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ def __init__(
dofs_motion_vel,
dofs_limit,
dofs_invweight,
dofs_frictionloss,
dofs_stiffness,
dofs_damping,
dofs_armature,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
"""
Expand Down
153 changes: 101 additions & 52 deletions genesis/engine/solvers/rigid/constraint_solver_decomp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -1295,9 +1307,32 @@ def func_ls_point_fn(
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]):
x = constraint_state.Jaref[i_c, i_b] + alpha * constraint_state.jv[i_c, 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
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 = x <= -rf
linear_pos = x >= rf
# active = (~linear_neg) & (~linear_pos)

if linear_neg or linear_pos:
qf = linear_neg * ti.Vector(
[f * (-0.5 * rf - constraint_state.Jaref[i_c, i_b]), -f * constraint_state.jv[i_c, i_b], 0.0]
)
qf += linear_pos * ti.Vector(
[f * (-0.5 * rf + constraint_state.Jaref[i_c, i_b]), f * constraint_state.jv[i_c, i_b], 0.0]
)
constraint_state.quad[i_c, _i0 + 0, i_b] = qf[0]
constraint_state.quad[i_c, _i0 + 1, i_b] = qf[1]
constraint_state.quad[i_c, _i0 + 2, i_b] = qf[2]

elif i_c >= nef:
active = x < 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
Expand Down Expand Up @@ -1681,9 +1716,23 @@ 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)
linear_neg = constraint_state.Jaref[i_c, i_b] <= -r * f
linear_pos = constraint_state.Jaref[i_c, i_b] >= r * f
constraint_state.active[i_c, i_b] = (~linear_neg) & (~linear_pos)
floss_force = linear_neg * f + linear_pos * -f
floss_cost_local = linear_neg * (-0.5 * r * f * f - f * constraint_state.Jaref[i_c, i_b])
floss_cost_local += linear_pos * (-0.5 * r * f * f + f * 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]
)

Expand Down
3 changes: 3 additions & 0 deletions genesis/engine/solvers/rigid/rigid_solver_decomp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -2513,6 +2514,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(),
Expand Down Expand Up @@ -2541,6 +2543,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]

Expand Down
6 changes: 6 additions & 0 deletions genesis/utils/array_class.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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_)),
Expand All @@ -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_)),
Expand Down Expand Up @@ -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
Expand All @@ -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),
Expand Down
4 changes: 2 additions & 2 deletions genesis/utils/geom.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -369,7 +369,7 @@ def ti_orthogonals(a):


@ti.func
def imp_aref(params, neg_penetration, vel, pos):
def imp_aref(params, neg_penetration, vel, pos_aref):
timeconst, dampratio, dmin, dmax, width, mid, power = params

imp_x = ti.abs(neg_penetration) / width
Expand All @@ -384,7 +384,7 @@ def imp_aref(params, neg_penetration, vel, pos):
b = 2.0 / (dmax * timeconst)
k = 1.0 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio)

aref = -b * vel - k * imp * pos
aref = -b * vel - k * imp * pos_aref

return imp, aref

Expand Down
Loading