diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 4d1dbfe9..91982152 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -504,7 +504,10 @@ def create_nmodel_batched_array(mjm_array, dtype): 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_stiffness=create_nmodel_batched_array(mjm.tendon_stiffness, dtype=float), + tendon_damping=create_nmodel_batched_array(mjm.tendon_damping, dtype=float), tendon_frictionloss=create_nmodel_batched_array(mjm.tendon_frictionloss, dtype=float), + tendon_lengthspring=create_nmodel_batched_array(mjm.tendon_lengthspring, dtype=wp.vec2), 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/passive.py b/mujoco_warp/_src/passive.py index 415bc694..9dc18a3c 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -27,22 +27,27 @@ @wp.kernel -def _spring_passive( +def _spring_damper_dof_passive( # Model: qpos_spring: wp.array2d(dtype=float), jnt_type: wp.array(dtype=int), jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), jnt_stiffness: wp.array2d(dtype=float), + dof_damping: wp.array2d(dtype=float), # Data in: qpos_in: wp.array2d(dtype=float), + qvel_in: wp.array2d(dtype=float), # Data out: qfrc_spring_out: wp.array2d(dtype=float), + qfrc_damper_out: wp.array2d(dtype=float), ): worldid, jntid = wp.tid() stiffness = jnt_stiffness[worldid, jntid] + dofid = jnt_dofadr[jntid] + damping = dof_damping[worldid, dofid] - if stiffness == 0.0: + if stiffness == 0.0 and damping == 0.0: return dofid = jnt_dofadr[jntid] @@ -50,6 +55,7 @@ def _spring_passive( qposid = jnt_qposadr[jntid] if jnttype == wp.static(JointType.FREE.value): + # spring dif = wp.vec3( qpos_in[worldid, qposid + 0] - qpos_spring[worldid, qposid + 0], qpos_in[worldid, qposid + 1] - qpos_spring[worldid, qposid + 1], @@ -75,7 +81,16 @@ def _spring_passive( qfrc_spring_out[worldid, dofid + 3] = -stiffness * dif[0] qfrc_spring_out[worldid, dofid + 4] = -stiffness * dif[1] qfrc_spring_out[worldid, dofid + 5] = -stiffness * dif[2] + + # damper + qfrc_damper_out[worldid, dofid + 0] = -damping * qvel_in[worldid, dofid + 0] + qfrc_damper_out[worldid, dofid + 1] = -damping * qvel_in[worldid, dofid + 1] + qfrc_damper_out[worldid, dofid + 2] = -damping * qvel_in[worldid, dofid + 2] + qfrc_damper_out[worldid, dofid + 3] = -damping * qvel_in[worldid, dofid + 3] + qfrc_damper_out[worldid, dofid + 4] = -damping * qvel_in[worldid, dofid + 4] + qfrc_damper_out[worldid, dofid + 5] = -damping * qvel_in[worldid, dofid + 5] elif jnttype == wp.static(JointType.BALL.value): + # spring rot = wp.quat( qpos_in[worldid, qposid + 0], qpos_in[worldid, qposid + 1], @@ -93,10 +108,67 @@ def _spring_passive( qfrc_spring_out[worldid, dofid + 0] = -stiffness * dif[0] qfrc_spring_out[worldid, dofid + 1] = -stiffness * dif[1] qfrc_spring_out[worldid, dofid + 2] = -stiffness * dif[2] + + # damper + qfrc_damper_out[worldid, dofid + 0] = -damping * qvel_in[worldid, dofid + 0] + qfrc_damper_out[worldid, dofid + 1] = -damping * qvel_in[worldid, dofid + 1] + qfrc_damper_out[worldid, dofid + 2] = -damping * qvel_in[worldid, dofid + 2] else: # mjJNT_SLIDE, mjJNT_HINGE fdif = qpos_in[worldid, qposid] - qpos_spring[worldid, qposid] qfrc_spring_out[worldid, dofid] = -stiffness * fdif + # damper + qfrc_damper_out[worldid, dofid] = -damping * qvel_in[worldid, dofid] + + +@wp.kernel +def _spring_damper_tendon_passive( + # Model: + tendon_stiffness: wp.array2d(dtype=float), + tendon_damping: wp.array2d(dtype=float), + tendon_lengthspring: wp.array2d(dtype=wp.vec2), + # Data in: + ten_velocity_in: wp.array2d(dtype=float), + ten_length_in: wp.array2d(dtype=float), + ten_J_in: wp.array3d(dtype=float), + # Data out: + qfrc_spring_out: wp.array2d(dtype=float), + qfrc_damper_out: wp.array2d(dtype=float), +): + worldid, tenid, dofid = wp.tid() + + stiffness = tendon_stiffness[worldid, tenid] + damping = tendon_damping[worldid, tenid] + + if stiffness == 0.0 and damping == 0.0: + return + + J = ten_J_in[worldid, tenid, dofid] + + if stiffness: + # compute spring force along tendon + length = ten_length_in[worldid, tenid] + lengthspring = tendon_lengthspring[worldid, tenid] + lower = lengthspring[0] + upper = lengthspring[1] + + if length > upper: + frc_spring = stiffness * (upper - length) + elif length < lower: + frc_spring = stiffness * (lower - length) + else: + frc_spring = 0.0 + + # transform to joint torque + wp.atomic_add(qfrc_spring_out[worldid], dofid, J * frc_spring) + + if damping: + # compute damper linear force along tendon + frc_damper = -damping * ten_velocity_in[worldid, tenid] + + # transform to joint torque + wp.atomic_add(qfrc_damper_out[worldid], dofid, J * frc_damper) + @wp.kernel def _gravity_force( @@ -251,29 +323,20 @@ def _qfrc_passive( # Model: jnt_actgravcomp: wp.array(dtype=int), dof_jntid: wp.array(dtype=int), - dof_damping: wp.array2d(dtype=float), # Data in: - qvel_in: wp.array2d(dtype=float), qfrc_spring_in: wp.array2d(dtype=float), + qfrc_damper_in: wp.array2d(dtype=float), qfrc_gravcomp_in: wp.array2d(dtype=float), qfrc_fluid_in: wp.array2d(dtype=float), # In: gravcomp: bool, fluid: bool, # Data out: - qfrc_damper_out: wp.array2d(dtype=float), qfrc_passive_out: wp.array2d(dtype=float), ): worldid, dofid = wp.tid() - - # spring qfrc_passive = qfrc_spring_in[worldid, dofid] - - # damper - qfrc_damper = -dof_damping[worldid, dofid] * qvel_in[worldid, dofid] - qfrc_damper_out[worldid, dofid] = qfrc_damper - - qfrc_passive += qfrc_damper + qfrc_passive += qfrc_damper_in[worldid, dofid] # add gravcomp unless added by actuators if gravcomp and not jnt_actgravcomp[dof_jntid[dofid]]: @@ -371,12 +434,41 @@ def passive(m: Model, d: Data): return d.qfrc_spring.zero_() + d.qfrc_damper.zero_() wp.launch( - _spring_passive, + _spring_damper_dof_passive, dim=(d.nworld, m.njnt), - inputs=[m.qpos_spring, m.jnt_type, m.jnt_qposadr, m.jnt_dofadr, m.jnt_stiffness, d.qpos], - outputs=[d.qfrc_spring], + inputs=[ + m.qpos_spring, + m.jnt_type, + m.jnt_qposadr, + m.jnt_dofadr, + m.jnt_stiffness, + m.dof_damping, + d.qpos, + d.qvel, + ], + outputs=[d.qfrc_spring, d.qfrc_damper], ) + + if m.ntendon: + wp.launch( + _spring_damper_tendon_passive, + dim=(d.nworld, m.ntendon, m.nv), + inputs=[ + m.tendon_stiffness, + m.tendon_damping, + m.tendon_lengthspring, + d.ten_velocity, + d.ten_length, + d.ten_J, + ], + outputs=[ + d.qfrc_spring, + d.qfrc_damper, + ], + ) + wp.launch( _flex_elasticity, dim=(d.nworld, m.nflexelem), @@ -428,16 +520,8 @@ def passive(m: Model, d: Data): wp.launch( _qfrc_passive, dim=(d.nworld, m.nv), - inputs=[ - m.jnt_actgravcomp, - m.dof_jntid, - m.dof_damping, - d.qvel, - d.qfrc_spring, - d.qfrc_gravcomp, - d.qfrc_fluid, - gravcomp, - fluid, + inputs=[m.jnt_actgravcomp, m.dof_jntid, d.qfrc_spring, d.qfrc_damper, d.qfrc_gravcomp, d.qfrc_fluid, gravcomp, fluid], + outputs=[ + d.qfrc_passive, ], - outputs=[d.qfrc_damper, d.qfrc_passive], ) diff --git a/mujoco_warp/_src/passive_test.py b/mujoco_warp/_src/passive_test.py index cac954ad..57e553f1 100644 --- a/mujoco_warp/_src/passive_test.py +++ b/mujoco_warp/_src/passive_test.py @@ -42,7 +42,7 @@ def test_passive(self, gravity): # TODO(taylorhowell): remove qpos0=True once tendon spring dampers are implemented _, mjd, m, d = test_util.fixture("pendula.xml", gravity=gravity, qpos0=True) - for arr in (d.qfrc_spring, d.qfrc_damper, d.qfrc_passive): + for arr in (d.qfrc_spring, d.qfrc_damper, d.qfrc_gravcomp, d.qfrc_passive): arr.zero_() mjwarp.passive(m, d) diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 1bc41ac3..af537b3d 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -773,7 +773,10 @@ class Model: 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_stiffness: stiffness coefficient (nworld, ntendon,) + tendon_damping: damping coefficient (nworld, ntendon,) tendon_frictionloss: loss due to friction (nworld, ntendon,) + tendon_lengthspring: spring resting length range (nworld, ntendon, 2) tendon_length0: tendon length in qpos0 (nworld, ntendon,) tendon_invweight0: inv. weight in qpos0 (nworld, ntendon,) wrap_objid: object id: geom, site, joint (nwrap,) @@ -1014,14 +1017,16 @@ 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_stiffness: wp.array2d(dtype=float) + tendon_damping: wp.array2d(dtype=float) tendon_frictionloss: wp.array2d(dtype=float) + tendon_lengthspring: wp.array2d(dtype=wp.vec2) 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 b65391f6..f451ecf7 100644 --- a/mujoco_warp/test_data/constraints.xml +++ b/mujoco_warp/test_data/constraints.xml @@ -78,7 +78,7 @@ - +