chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove mjDSBL_PASSIVE, split it into mjDSBL_SPRING and mjDSBL_DAMPING.

Browse Source 
PiperOrigin-RevId: 800870384
Change-Id: I273e0ceec090901e2620febb67e57914928e94bb
This commit is contained in:
Yuval Tassa
2025-08-29 06:05:18 -07:00
committed by Copybara-Service
parent 0eea5d6dbd
commit b092563cab
39 changed files with 455 additions and 289 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/derivative.py
+3 -1
View File
@@ -55,13 +55,15 @@ def deriv_smooth_vel(m: Model, d: Data) -> Optional[jax.Array]:
)
# qDeriv += d qfrc_passive / d qvel
if not m.opt.disableflags & DisableBit.PASSIVE:
if not m.opt.disableflags & DisableBit.DAMPER:
if qderiv is None:
qderiv = -jp.diag(m.dof_damping)
else:
qderiv -= jp.diag(m.dof_damping)
if m.ntendon:
qderiv -= d._impl.ten_J.T @ jp.diag(m.tendon_damping) @ d._impl.ten_J
if not m.opt.disableflags & (DisableBit.DAMPER | DisableBit.SPRING):
# TODO(robotics-simulation): fluid drag model
if m.opt._impl.has_fluid_params: # pytype: disable=attribute-error
raise NotImplementedError('fluid drag not supported for implicitfast')
mjx/mujoco/mjx/_src/passive.py
+30 -17
View File
@@ -60,25 +60,38 @@ def _spring_damper(m: Model, d: Data) -> jax.Array:
return jp.concatenate(qfrcs)
# dof-level springs
qfrc = scan.flat(
m,
fn,
'jjqq',
'v',
m.jnt_type,
m.jnt_stiffness,
m.qpos_spring,
d.qpos,
)
qfrc = jp.zeros(m.nv)
if not m.opt.disableflags & DisableBit.SPRING:
qfrc = scan.flat(
m,
fn,
'jjqq',
'v',
m.jnt_type,
m.jnt_stiffness,
m.qpos_spring,
d.qpos,
)
# dof-level dampers
qfrc -= m.dof_damping * d.qvel
if not m.opt.disableflags & DisableBit.DAMPER:
qfrc -= m.dof_damping * d.qvel
# tendon-level springs
if not m.opt.disableflags & DisableBit.SPRING:
below, above = m.tendon_lengthspring.T - d.ten_length
frc_spring = jp.where(below > 0, m.tendon_stiffness * below, 0)
frc_spring = jp.where(above < 0, m.tendon_stiffness * above, frc_spring)
else:
frc_spring = jp.zeros(m.ntendon)
# tendon-level dampers
frc_damper = (
-m.tendon_damping * d._impl.ten_velocity
if not m.opt.disableflags & DisableBit.DAMPER
else jp.zeros(m.ntendon)
)
# tendon-level spring-dampers
below, above = m.tendon_lengthspring.T - d.ten_length
frc_spring = jp.where(below > 0, m.tendon_stiffness * below, 0)
frc_spring = jp.where(above < 0, m.tendon_stiffness * above, frc_spring)
frc_damper = -m.tendon_damping * d._impl.ten_velocity
qfrc += d._impl.ten_J.T @ (frc_spring + frc_damper)
return qfrc
@@ -123,7 +136,7 @@ def passive(m: Model, d: Data) -> Data:
):
raise ValueError('passive requires JAX backend implementation.')
if m.opt.disableflags & DisableBit.PASSIVE:
if m.opt.disableflags & (DisableBit.SPRING | DisableBit.DAMPER):
return d.replace(qfrc_passive=jp.zeros(m.nv), qfrc_gravcomp=jp.zeros(m.nv))
qfrc_passive = _spring_damper(m, d)
mjx/mujoco/mjx/_src/passive_test.py
+3 -1
View File
@@ -76,7 +76,9 @@ class PassiveTest(absltest.TestCase):
_assert_attr_eq(d, dx, 'qfrc_gravcomp')
# test disable passive
mx = mx.tree_replace({'opt.disableflags': mjx.DisableBit.PASSIVE})
mx = mx.tree_replace(
{'opt.disableflags': mjx.DisableBit.SPRING | mjx.DisableBit.DAMPER}
)
dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
np.testing.assert_allclose(dx.qfrc_passive, 0)
mjx/mujoco/mjx/_src/types.py
+4 -2
View File
@@ -53,7 +53,8 @@ class DisableBit(enum.IntFlag):
FRICTIONLOSS: joint and tendon frictionloss constraints
LIMIT: joint and tendon limit constraints
CONTACT: contact constraints
PASSIVE: passive forces
SPRING: passive spring forces
DAMPER: passive damper forces
GRAVITY: gravitational forces
CLAMPCTRL: clamp control to specified range
WARMSTART: warmstart constraint solver
@@ -67,7 +68,8 @@ class DisableBit(enum.IntFlag):
FRICTIONLOSS = mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS
LIMIT = mujoco.mjtDisableBit.mjDSBL_LIMIT
CONTACT = mujoco.mjtDisableBit.mjDSBL_CONTACT
PASSIVE = mujoco.mjtDisableBit.mjDSBL_PASSIVE
SPRING = mujoco.mjtDisableBit.mjDSBL_SPRING
DAMPER = mujoco.mjtDisableBit.mjDSBL_DAMPER
GRAVITY = mujoco.mjtDisableBit.mjDSBL_GRAVITY
CLAMPCTRL = mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL
WARMSTART = mujoco.mjtDisableBit.mjDSBL_WARMSTART
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+1 -1
View File
@@ -140,7 +140,7 @@ def deriv_smooth_vel(m: Model, d: Data, flg_forward: bool = True):
Default is True.
"""
actuation_enabled = not (m.opt.disableflags & DisableBit.ACTUATION)
passive_enabled = not (m.opt.disableflags & DisableBit.PASSIVE)
passive_enabled = not (m.opt.disableflags & (DisableBit.SPRING | DisableBit.DAMPER))
qMi = m.qM_fullm_i if m.opt.is_sparse else m.dof_tri_row
qMj = m.qM_fullm_j if m.opt.is_sparse else m.dof_tri_col
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+1 -1
View File
@@ -523,7 +523,7 @@ def implicit(m: Model, d: Data):
"""Integrates fully implicit in velocity."""
# compile-time constants
passive_enabled = not m.opt.disableflags & DisableBit.PASSIVE.value
passive_enabled = not m.opt.disableflags & (DisableBit.SPRING.value | DisableBit.DAMPER.value)
actuation_enabled = (not m.opt.disableflags & DisableBit.ACTUATION.value) and m.actuator_affine_bias_gain
if passive_enabled or actuation_enabled:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward_test.py
Vendored
+2 -1
View File
@@ -194,7 +194,8 @@ class ForwardTest(parameterized.TestCase):
"pendula.xml",
integrator=IntegratorType.IMPLICITFAST,
actuation=actuation,
passive=passive,
spring=passive,
damper=passive,
sparse=sparse,
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive.py
Vendored
+1 -1
View File
@@ -505,7 +505,7 @@ def _flex_bending(
def passive(m: Model, d: Data):
"""Adds all passive forces."""
if m.opt.disableflags & DisableBit.PASSIVE:
if m.opt.disableflags & (DisableBit.SPRING | DisableBit.DAMPER):
d.qfrc_spring.zero_()
d.qfrc_damper.zero_()
d.qfrc_gravcomp.zero_()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/passive_test.py
Vendored
+1 -1
View File
@@ -39,7 +39,7 @@ class PassiveTest(parameterized.TestCase):
@parameterized.product(passive=[True, False], gravity=[True, False])
def test_passive(self, passive, gravity):
"""Tests passive."""
_, mjd, m, d = test_util.fixture("pendula.xml", passive=passive, gravity=gravity, kick=True, applied=True)
_, mjd, m, d = test_util.fixture("pendula.xml", spring=passive, damper=passive, gravity=gravity, kick=True, applied=True)
for arr in (d.qfrc_spring, d.qfrc_damper, d.qfrc_gravcomp, d.qfrc_passive):
arr.zero_()
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+1 -1
View File
@@ -2370,7 +2370,7 @@ def energy_pos(m: Model, d: Data):
_energy_pos_gravity, dim=(d.nworld, m.nbody - 1), inputs=[m.opt.gravity, m.body_mass, d.xipos], outputs=[d.energy]
)
if not m.opt.disableflags & DisableBit.PASSIVE:
if not m.opt.disableflags & DisableBit.SPRING:
# add joint-level springs
wp.launch(
_energy_pos_passive_joint,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py
Vendored
+6 -3
View File
@@ -46,7 +46,8 @@ def fixture(
contact: bool = True,
constraint: bool = True,
equality: bool = True,
passive: bool = True,
spring: bool = True,
damper: bool = True,
gravity: bool = True,
clampctrl: bool = True,
filterparent: bool = True,
@@ -87,8 +88,10 @@ def fixture(
mjm.opt.disableflags |= DisableBit.CONSTRAINT
if not equality:
mjm.opt.disableflags |= DisableBit.EQUALITY
if not passive:
mjm.opt.disableflags |= DisableBit.PASSIVE
if not spring:
mjm.opt.disableflags |= DisableBit.SPRING
if not damper:
mjm.opt.disableflags |= DisableBit.DAMPER
if not gravity:
mjm.opt.disableflags |= DisableBit.GRAVITY
if not clampctrl:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+4 -2
View File
@@ -129,7 +129,8 @@ class DisableBit(enum.IntFlag):
FRICTIONLOSS: joint and tendon frictionloss constraints
LIMIT: joint and tendon limit constraints
CONTACT: contact constraints
PASSIVE: passive forces
SPRING: passive spring forces
DAMPER: passive damper forces
GRAVITY: gravitational forces
CLAMPCTRL: clamp control to specified range
ACTUATION: apply actuation forces
@@ -144,7 +145,8 @@ class DisableBit(enum.IntFlag):
FRICTIONLOSS = mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS
LIMIT = mujoco.mjtDisableBit.mjDSBL_LIMIT
CONTACT = mujoco.mjtDisableBit.mjDSBL_CONTACT
PASSIVE = mujoco.mjtDisableBit.mjDSBL_PASSIVE
SPRING = mujoco.mjtDisableBit.mjDSBL_SPRING
DAMPER = mujoco.mjtDisableBit.mjDSBL_DAMPER
GRAVITY = mujoco.mjtDisableBit.mjDSBL_GRAVITY
CLAMPCTRL = mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL
WARMSTART = mujoco.mjtDisableBit.mjDSBL_WARMSTART