diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index dfec215c..2f625e2c 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -437,7 +437,7 @@ def _efc_equality_tendon( @wp.kernel -def _efc_friction( +def _efc_friction_dof( # Model: opt_timestep: float, dof_invweight0: wp.array2d(dtype=float), @@ -461,7 +461,6 @@ def _efc_friction( efc_aref_out: wp.array(dtype=float), efc_frictionloss_out: wp.array(dtype=float), ): - # TODO(team): tendon worldid, dofid = wp.tid() if dof_frictionloss[worldid, dofid] <= 0.0: @@ -500,6 +499,72 @@ def _efc_friction( ) +@wp.kernel +def _efc_friction_tendon( + # Model: + nv: int, + opt_timestep: float, + tendon_solref_fri: wp.array2d(dtype=wp.vec2), + tendon_solimp_fri: wp.array2d(dtype=vec5), + tendon_frictionloss: wp.array2d(dtype=float), + tendon_invweight0: wp.array2d(dtype=float), + # Data in: + qvel_in: wp.array2d(dtype=float), + ten_J_in: wp.array3d(dtype=float), + # In: + refsafe_in: int, + # Data out: + nf_out: wp.array(dtype=int), + nefc_out: wp.array(dtype=int), + efc_worldid_out: wp.array(dtype=int), + efc_id_out: wp.array(dtype=int), + efc_J_out: wp.array2d(dtype=float), + efc_pos_out: wp.array(dtype=float), + efc_margin_out: wp.array(dtype=float), + efc_D_out: wp.array(dtype=float), + efc_aref_out: wp.array(dtype=float), + efc_frictionloss_out: wp.array(dtype=float), +): + worldid, tenid = wp.tid() + + frictionloss = tendon_frictionloss[worldid, tenid] + if frictionloss <= 0.0: + return + + efcid = wp.atomic_add(nefc_out, 0, 1) + wp.atomic_add(nf_out, 0, 1) + efc_worldid_out[efcid] = worldid + + Jqvel = float(0.0) + + # TODO(team): parallelize + for i in range(nv): + J = ten_J_in[worldid, tenid, i] + efc_J_out[efcid, i] = J + Jqvel += J * qvel_in[worldid, i] + + _update_efc_row( + opt_timestep, + refsafe_in, + efcid, + 0.0, + 0.0, + tendon_invweight0[worldid, tenid], + tendon_solref_fri[worldid, tenid], + tendon_solimp_fri[worldid, tenid], + 0.0, + Jqvel, + frictionloss, + tenid, + efc_id_out, + efc_pos_out, + efc_margin_out, + efc_D_out, + efc_aref_out, + efc_frictionloss_out, + ) + + @wp.kernel def _efc_equality_weld( # Model: @@ -1477,7 +1542,7 @@ def make_constraint(m: types.Model, d: types.Data): if not (m.opt.disableflags & types.DisableBit.FRICTIONLOSS.value): wp.launch( - _efc_friction, + _efc_friction_dof, dim=(d.nworld, m.nv), inputs=[ m.opt.timestep, @@ -1503,6 +1568,34 @@ def make_constraint(m: types.Model, d: types.Data): ], ) + wp.launch( + _efc_friction_tendon, + dim=(d.nworld, m.ntendon), + inputs=[ + m.nv, + m.opt.timestep, + m.tendon_solref_fri, + m.tendon_solimp_fri, + m.tendon_frictionloss, + m.tendon_invweight0, + d.qvel, + d.ten_J, + refsafe, + ], + outputs=[ + d.nf, + d.nefc, + d.efc.worldid, + d.efc.id, + d.efc.J, + d.efc.pos, + d.efc.margin, + d.efc.D, + d.efc.aref, + d.efc.frictionloss, + ], + ) + # limit if not (m.opt.disableflags & types.DisableBit.LIMIT.value): limit_ball = m.jnt_limited_ball_adr.size > 0 diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 8dea42bb..4d1dbfe9 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -44,9 +44,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: if mjm.nflex > 1: raise NotImplementedError("Only one flex is unsupported.") - if mjm.tendon_frictionloss.any(): - raise NotImplementedError("Tendon frictionloss is unsupported.") - if mjm.geom_fluid.any(): raise NotImplementedError("Ellipsoid fluid model not implemented.") @@ -503,8 +500,11 @@ def create_nmodel_batched_array(mjm_array, dtype): tendon_limited_adr=wp.array(np.nonzero(mjm.tendon_limited)[0], dtype=wp.int32, ndim=1), tendon_solref_lim=create_nmodel_batched_array(mjm.tendon_solref_lim, dtype=wp.vec2f), tendon_solimp_lim=create_nmodel_batched_array(mjm.tendon_solimp_lim, dtype=types.vec5), + tendon_solref_fri=create_nmodel_batched_array(mjm.tendon_solref_fri, dtype=wp.vec2f), + tendon_solimp_fri=create_nmodel_batched_array(mjm.tendon_solimp_fri, dtype=types.vec5), tendon_range=create_nmodel_batched_array(mjm.tendon_range, dtype=wp.vec2f), tendon_margin=create_nmodel_batched_array(mjm.tendon_margin, dtype=float), + tendon_frictionloss=create_nmodel_batched_array(mjm.tendon_frictionloss, dtype=float), tendon_length0=create_nmodel_batched_array(mjm.tendon_length0, dtype=float), tendon_invweight0=create_nmodel_batched_array(mjm.tendon_invweight0, dtype=float), wrap_objid=wp.array(mjm.wrap_objid, dtype=int), diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 2c6d1acb..1bc41ac3 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -769,8 +769,11 @@ class Model: tendon_limited_adr: addresses for limited tendons (<=ntendon,) tendon_solref_lim: constraint solver reference: limit (nworld, ntendon, mjNREF) tendon_solimp_lim: constraint solver impedance: limit (nworld, ntendon, mjNIMP) + tendon_solref_fri: constraint solver reference: friction (nworld, ntendon, mjNREF) + tendon_solimp_fri: constraint solver impedance: friction (nworld, ntendon, mjNIMP) tendon_range: tendon length limits (nworld, ntendon, 2) tendon_margin: min distance for limit detection (nworld, ntendon,) + tendon_frictionloss: loss due to friction (nworld, ntendon,) tendon_length0: tendon length in qpos0 (nworld, ntendon,) tendon_invweight0: inv. weight in qpos0 (nworld, ntendon,) wrap_objid: object id: geom, site, joint (nwrap,) @@ -1011,10 +1014,14 @@ class Model: tendon_num: wp.array(dtype=int) tendon_limited: wp.array(dtype=int) tendon_limited_adr: wp.array(dtype=int) + tendon_solref_lim: wp.array2d(dtype=wp.vec2) tendon_solimp_lim: wp.array2d(dtype=vec5) + tendon_solref_fri: wp.array2d(dtype=wp.vec2) + tendon_solimp_fri: wp.array2d(dtype=vec5) tendon_range: wp.array2d(dtype=wp.vec2) tendon_margin: wp.array2d(dtype=float) + tendon_frictionloss: wp.array2d(dtype=float) tendon_length0: wp.array2d(dtype=float) tendon_invweight0: wp.array2d(dtype=float) wrap_objid: wp.array(dtype=int) diff --git a/mujoco_warp/test_data/constraints.xml b/mujoco_warp/test_data/constraints.xml index 4e8b4760..b65391f6 100644 --- a/mujoco_warp/test_data/constraints.xml +++ b/mujoco_warp/test_data/constraints.xml @@ -77,6 +77,13 @@ + + + + + + +