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

Add gravcomp.

Browse Source 
PiperOrigin-RevId: 627791303
Change-Id: I87c9254967eb492d0c38153ff41d7458ffb69136
This commit is contained in:
Baruch Tabanpour
2024-04-24 11:18:35 -07:00
committed by Copybara-Service
parent 68e33f4c4e
commit e9709900b4
10 changed files with 152 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/forward.py
+5 -4
View File
@@ -174,16 +174,17 @@ def fwd_actuation(m: Model, d: Data) -> Data:
qfrc_actuator = d.actuator_moment.T @ force
if m.ngravcomp:
# actuator-level gravity compensation, skip if added as passive force
qfrc_actuator += d.qfrc_gravcomp * m.jnt_actgravcomp[m.dof_jntid]
# clamp qfrc_actuator
actfrcrange = jp.where(
m.jnt_actfrclimited[:, None],
m.jnt_actfrcrange,
jp.array([-jp.inf, jp.inf]),
)
ids = sum(
([i] * JointType(j).dof_width() for i, j in enumerate(m.jnt_type)), []
)
actfrcrange = jp.take(actfrcrange, jp.array(ids), axis=0)
actfrcrange = actfrcrange[m.dof_jntid]
qfrc_actuator = jp.clip(qfrc_actuator, actfrcrange[:, 0], actfrcrange[:, 1])
d = d.replace(act_dot=act_dot, qfrc_actuator=qfrc_actuator)
mjx/mujoco/mjx/_src/io.py
+1 -3
View File
@@ -68,9 +68,6 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
if m.ntendon:
raise NotImplementedError('tendons are not supported')
if m.body_gravcomp.any():
raise NotImplementedError('gravcomp is not supported')
for g1, g2, ip in collision_driver.geom_pairs(m):
t1, t2 = m.geom_type[[g1, g2]]
# check collision function exists for type pair
@@ -201,6 +198,7 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
cvel=zero_nbody_6,
cdof_dot=zero_nv_6,
qfrc_bias=zero_nv,
qfrc_gravcomp=zero_nv,
qfrc_passive=zero_nv,
efc_aref=zero_nefc,
qfrc_actuator=zero_nv,
mjx/mujoco/mjx/_src/io_test.py
-12
View File
@@ -173,18 +173,6 @@ class ModelIOTest(parameterized.TestCase):
</tendon>
</mujoco>"""))
def test_gravcomp_not_implemented(self):
with self.assertRaises(NotImplementedError):
mjx.put_model(mujoco.MjModel.from_xml_string("""
<mujoco>
<worldbody>
<body gravcomp="1">
<freejoint/>
<geom size="0.05"/>
</body>
</worldbody>
</mujoco>"""))
def test_cylinder_not_implemented(self):
with self.assertRaises(NotImplementedError):
mjx.put_model(mujoco.MjModel.from_xml_string("""
mjx/mujoco/mjx/_src/passive.py
+94 -68
View File
@@ -29,6 +29,100 @@ from mujoco.mjx._src.types import Model
# pylint: enable=g-importing-member
def _spring_damper(m: Model, d: Data) -> jax.Array:
"""Applies joint level spring and damping forces."""
def fn(jnt_typs, stiffness, qpos_spring, qpos):
qpos_i = 0
qfrcs = []
for i in range(len(jnt_typs)):
jnt_typ = JointType(jnt_typs[i])
q = qpos[qpos_i : qpos_i + jnt_typ.qpos_width()]
qs = qpos_spring[qpos_i : qpos_i + jnt_typ.qpos_width()]
qfrc = jp.zeros(jnt_typ.dof_width())
if jnt_typ == JointType.FREE:
qfrc = qfrc.at[:3].set(-stiffness[i] * (q[:3] - qs[:3]))
qfrc = qfrc.at[3:6].set(-stiffness[i] * math.quat_sub(q[3:7], qs[3:7]))
elif jnt_typ == JointType.BALL:
qfrc = -stiffness[i] * math.quat_sub(q, qs)
elif jnt_typ in (
JointType.SLIDE,
JointType.HINGE,
):
qfrc = -stiffness[i] * (q - qs)
else:
raise RuntimeError(f'unrecognized joint type: {jnt_typ}')
qfrcs.append(qfrc)
qpos_i += jnt_typ.qpos_width()
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,
)
# dof-level dampers
qfrc -= m.dof_damping * d.qvel
return qfrc
def _gravcomp(m: Model, d: Data) -> jax.Array:
"""Applies body-level gravity compensation."""
force = -m.opt.gravity * (m.body_mass * m.body_gravcomp)[:, None]
apply_f = lambda f, pos, body_id: support.jac(m, d, pos, body_id)[0] @ f
qfrc = jax.vmap(apply_f)(force, d.xipos, jp.arange(m.nbody)).sum(axis=0)
return qfrc
def _fluid(m: Model, d: Data) -> jax.Array:
"""Applies body-level viscosity, lift and drag."""
force, torque = jax.vmap(
_inertia_box_fluid_model, in_axes=(None, 0, 0, 0, 0, 0, 0)
)(
m,
m.body_inertia,
m.body_mass,
d.subtree_com[jp.array(m.body_rootid)],
d.xipos,
d.ximat,
d.cvel,
)
qfrc = jax.vmap(support.apply_ft, in_axes=(None, None, 0, 0, 0, 0))(
m, d, force, torque, d.xipos, jp.arange(m.nbody)
)
return jp.sum(qfrc, axis=0)
def passive(m: Model, d: Data) -> Data:
"""Adds all passive forces."""
if m.opt.disableflags & DisableBit.PASSIVE:
return d.replace(qfrc_passive=jp.zeros(m.nv), qfrc_gravcomp=jp.zeros(m.nv))
qfrc_passive = _spring_damper(m, d)
qfrc_gravcomp = jp.zeros(m.nv)
if m.ngravcomp and not m.opt.disableflags & DisableBit.GRAVITY:
qfrc_gravcomp = _gravcomp(m, d)
# add gravcomp unless added via actuators
qfrc_passive += qfrc_gravcomp * (1 - m.jnt_actgravcomp[m.dof_jntid])
if m.opt.has_fluid_params:
qfrc_passive += _fluid(m, d)
d = d.replace(qfrc_passive=qfrc_passive, qfrc_gravcomp=qfrc_gravcomp)
return d
def _inertia_box_fluid_model(
m: Model,
inertia: jax.Array,
@@ -71,71 +165,3 @@ def _inertia_box_fluid_model(
force, torque = ximat @ lfrc_vel, ximat @ lfrc_ang
return force, torque
def passive(m: Model, d: Data) -> Data:
"""Adds all passive forces."""
if m.opt.disableflags & DisableBit.PASSIVE:
return d.replace(qfrc_passive=jp.zeros(m.nv))
# joint-level springs
def fn(jnt_typs, stiffness, qpos_spring, qpos):
qpos_i = 0
qfrcs = []
for i in range(len(jnt_typs)):
jnt_typ = JointType(jnt_typs[i])
q = qpos[qpos_i : qpos_i + jnt_typ.qpos_width()]
qs = qpos_spring[qpos_i : qpos_i + jnt_typ.qpos_width()]
qfrc = jp.zeros(jnt_typ.dof_width())
if jnt_typ == JointType.FREE:
qfrc = qfrc.at[:3].set(-stiffness[i] * (q[:3] - qs[:3]))
qfrc = qfrc.at[3:6].set(-stiffness[i] * math.quat_sub(q[3:7], qs[3:7]))
elif jnt_typ == JointType.BALL:
qfrc = -stiffness[i] * math.quat_sub(q, qs)
elif jnt_typ in (
JointType.SLIDE,
JointType.HINGE,
):
qfrc = -stiffness[i] * (q - qs)
else:
raise RuntimeError(f'unrecognized joint type: {jnt_typ}')
qfrcs.append(qfrc)
qpos_i += jnt_typ.qpos_width()
return jp.concatenate(qfrcs)
qfrc_passive = scan.flat(
m,
fn,
'jjqq',
'v',
m.jnt_type,
m.jnt_stiffness,
m.qpos_spring,
d.qpos,
)
# dof-level dampers
qfrc_passive -= m.dof_damping * d.qvel
# TODO(robotics-simulation): body-level gravity compensation
# body-level viscosity, lift and drag
if m.opt.has_fluid_params:
force, torque = jax.vmap(
_inertia_box_fluid_model, in_axes=(None, 0, 0, 0, 0, 0, 0)
)(
m,
m.body_inertia,
m.body_mass,
d.subtree_com[jp.array(m.body_rootid)],
d.xipos,
d.ximat,
d.cvel,
)
qfrc_target = jax.vmap(support.apply_ft, in_axes=(None, None, 0, 0, 0, 0))(
m, d, force, torque, d.xipos, jp.arange(m.nbody)
)
qfrc_passive += jp.sum(qfrc_target, axis=0)
d = d.replace(qfrc_passive=qfrc_passive)
return d
mjx/mujoco/mjx/_src/passive_test.py
+5 -1
View File
@@ -42,13 +42,14 @@ class PassiveTest(absltest.TestCase):
m = test_util.load_test_file('pendula.xml')
d = mujoco.MjData(m)
# give the system a little kick to ensure we have non-identity rotations
d.ctrl = np.array([0.1, -0.1, 0.2, 0.3, -0.4])
d.ctrl = np.array([0.1, -0.1, 0.2, 0.3, -0.4, 0.5, -0.6, 0.1])
mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qfrc_passive')
_assert_attr_eq(d, dx, 'qfrc_gravcomp')
# test with fluid forces
m.opt.density = 0.01
@@ -56,18 +57,21 @@ class PassiveTest(absltest.TestCase):
mx = mjx.put_model(m)
dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qfrc_passive')
_assert_attr_eq(d, dx, 'qfrc_gravcomp')
m.opt.viscosity = 0.02
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qfrc_passive')
_assert_attr_eq(d, dx, 'qfrc_gravcomp')
m.opt.wind = np.array([0.03, 0.04, 0.05])
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qfrc_passive')
_assert_attr_eq(d, dx, 'qfrc_gravcomp')
# test disable passive
mx = mx.tree_replace({'opt.disableflags': mjx.DisableBit.PASSIVE})
mjx/mujoco/mjx/_src/test_util.py
+11 -1
View File
@@ -152,6 +152,7 @@ def _make_joint(joint_type: str, name: str) -> Dict[str, str]:
joint_attr['damping'] = '{:.2f}'.format(np.random.uniform() * 20)
joint_attr['stiffness'] = '{:.2f}'.format(np.random.uniform() * 20)
joint_attr['actuatorgravcomp'] = np.random.choice(['true', 'false'])
return joint_attr
@@ -334,7 +335,16 @@ def create_mjcf(
z_pos = np.random.uniform(low=-1, high=1) * 0.01 # small jitter
pos = f'{body_pos[0]:.3f} {body_pos[1]:.3f} {body_pos[2] + z_pos:.3f}'
n_bodies = len(list(mjcf.iter('body')))
child = ET.SubElement(body, 'body', {'pos': pos, 'name': f'body{n_bodies}'})
gravcomp = np.random.uniform() * p(50)
child = ET.SubElement(
body,
'body',
{
'pos': pos,
'name': f'body{n_bodies}',
'gravcomp': f'{gravcomp:.3f}',
},
)
ET.SubElement(child, 'site', {'name': f'site{n_bodies}'})
n_joints = len(list(mjcf.iter('joint')))
mjx/mujoco/mjx/_src/types.py
+10 -1
View File
@@ -326,6 +326,7 @@ class Model(PyTreeNode):
npair: number of predefined geom pairs
nexclude: number of excluded geom pairs
neq: number of equality constraints
ngravcomp: number of bodies with nonzero gravcomp
nnumeric: number of numeric custom fields
ntuple: number of tuple custom fields
nsensor: number of sensors
@@ -352,6 +353,7 @@ class Model(PyTreeNode):
body_mass: mass (nbody,)
body_subtreemass: mass of subtree starting at this body (nbody,)
body_inertia: diagonal inertia in ipos/iquat frame (nbody, 3)
body_gravcomp: antigravity force, units of body weight (nbody,)
body_invweight0: mean inv inert in qpos0 (trn, rot) (nbody, 2)
jnt_type: type of joint (mjtJoint) (njnt,)
jnt_qposadr: start addr in 'qpos' for joint's data (njnt,)
@@ -359,6 +361,8 @@ class Model(PyTreeNode):
jnt_bodyid: id of joint's body (njnt,)
jnt_group: group for visibility (njnt,)
jnt_limited: does joint have limits (njnt,)
jnt_actfrclimited: does joint have actuator force limits (njnt,)
jnt_actgravcomp: is gravcomp force applied via actuators (njnt,)
jnt_solref: constraint solver reference: limit (njnt, mjNREF)
jnt_solimp: constraint solver impedance: limit (njnt, mjNIMP)
jnt_pos: local anchor position (njnt, 3)
@@ -488,6 +492,7 @@ class Model(PyTreeNode):
npair: int
nexclude: int
neq: int
ngravcomp: int
nnumeric: int
nuserdata: int
ntuple: int
@@ -514,6 +519,7 @@ class Model(PyTreeNode):
body_mass: jax.Array
body_subtreemass: jax.Array
body_inertia: jax.Array
body_gravcomp: jax.Array
body_invweight0: jax.Array
jnt_type: np.ndarray
jnt_qposadr: np.ndarray
@@ -521,6 +527,7 @@ class Model(PyTreeNode):
jnt_bodyid: np.ndarray
jnt_limited: np.ndarray
jnt_actfrclimited: np.ndarray
jnt_actgravcomp: np.ndarray
jnt_solref: jax.Array
jnt_solimp: jax.Array
jnt_pos: jax.Array
@@ -671,7 +678,7 @@ class Contact(PyTreeNode):
class Data(PyTreeNode):
r"""Dynamic state that updates each step.
r"""\Dynamic state that updates each step.
Attributes:
ne: number of equality constraints
@@ -725,6 +732,7 @@ class Data(PyTreeNode):
cvel: com-based velocity [3D rot; 3D tran] (nbody, 6)
cdof_dot: time-derivative of cdof (nv, 6)
qfrc_bias: C(qpos,qvel) (nv,)
qfrc_gravcomp: passive gravity compensation force (nv,)
qfrc_passive: passive force (nv,)
efc_aref: reference pseudo-acceleration (nefc,)
qfrc_actuator: actuator force (nv,)
@@ -795,6 +803,7 @@ class Data(PyTreeNode):
cdof_dot: jax.Array
qfrc_bias: jax.Array
qfrc_passive: jax.Array
qfrc_gravcomp: jax.Array
efc_aref: jax.Array
# position, velcoity, control & acceleration dependent:
qfrc_actuator: jax.Array
mjx/mujoco/mjx/integration_test/smooth_test.py
+4 -4
View File
@@ -68,10 +68,10 @@ class TransmissionIntegrationTest(parameterized.TestCase):
mujoco.mj_transmission(m, d)
dx = transmission_jit_fn(mx, dx)
_assert_attr_eq(d, dx, 'actuator_length', seed, f'transmission{seed}')
_assert_attr_eq(
d, dx, 'actuator_moment', seed, f'transmission{seed}', atol=1e-4
)
for field in ['actuator_length', 'actuator_moment']:
_assert_attr_eq(
d, dx, field, seed, f'transmission{seed}', atol=1e-4
)
if __name__ == '__main__':
mjx/mujoco/mjx/test_data/pendula.xml
+17
View File
@@ -110,6 +110,20 @@
</body>
</body>
</body>
<!-- triple pendulum of hinges with gravcomp -->
<body pos="3.0 0 0" gravcomp="1">
<joint name="joint15" axis="0.1 0.2 0.3" type="hinge"/>
<geom/>
<body pos="0 0 -0.8" gravcomp="2">
<joint name="joint16" axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
<geom/>
<body pos="0 -0.7 0" gravcomp="3">
<joint name="joint17" axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30" actuatorgravcomp="true"/>
<geom/>
</body>
</body>
</body>
</worldbody>
<actuator>
@@ -118,5 +132,8 @@
<motor gear="0 0 300" joint="joint1" name="act3"/>
<motor gear="275" joint="joint2" name="act4"/>
<motor gear="275" joint="joint3" name="act5"/>
<motor gear="150" joint="joint15" name="act6"/>
<motor gear="150" joint="joint16" name="act7"/>
<motor gear="150" joint="joint17" name="act8"/>
</actuator>
</mujoco>