From 143196ee79dab87658f0a44c30b8512c31f16e64 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 9 Jul 2024 18:59:23 +0000 Subject: [PATCH 01/26] Bump zipp from 3.17.0 to 3.19.1 in /python Bumps [zipp](https://github.com/jaraco/zipp) from 3.17.0 to 3.19.1. - [Release notes](https://github.com/jaraco/zipp/releases) - [Changelog](https://github.com/jaraco/zipp/blob/main/NEWS.rst) - [Commits](https://github.com/jaraco/zipp/compare/v3.17.0...v3.19.1) --- updated-dependencies: - dependency-name: zipp dependency-type: direct:production ... Signed-off-by: dependabot[bot] --- python/build_requirements.txt | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/build_requirements.txt b/python/build_requirements.txt index e927a0bf..506ef975 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -77,8 +77,9 @@ importlib-resources==6.1.0 \ --hash=sha256:aa50258bbfa56d4e33fbd8aa3ef48ded10d1735f11532b8df95388cc6bdb7e83 typing_extensions==4.8.0 \ --hash=sha256:8f92fc8806f9a6b641eaa5318da32b44d401efaac0f6678c9bc448ba3605faa0 -zipp==3.17.0 \ - --hash=sha256:0e923e726174922dce09c53c59ad483ff7bbb8e572e00c7f7c46b88556409f31 +zipp==3.19.1 \ + --hash=sha256:2828e64edb5386ea6a52e7ba7cdb17bb30a73a858f5eb6eb93d8d36f5ea26091 \ + --hash=sha256:35427f6d5594f4acf82d25541438348c26736fa9b3afa2754bcd63cdb99d8e8f # Transitive dependencies of pytest attrs==23.1.0; platform_system == 'Windows' \ From 6acf40613d430d5c06433b431e33658cfea4a18b Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Sun, 18 Aug 2024 04:40:24 -0700 Subject: [PATCH 02/26] Add support for disabling MJX sensors. PiperOrigin-RevId: 664410169 Change-Id: Icb278e9d7a9a493d5a61ab04c1300d8e412d0e83 --- mjx/mujoco/mjx/_src/sensor.py | 12 ++++++++++++ mjx/mujoco/mjx/_src/sensor_test.py | 20 ++++++++++++++++++++ mjx/mujoco/mjx/_src/types.py | 4 +++- 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index 7dabeeb2..e644bcdc 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -20,6 +20,7 @@ import mujoco # pylint: disable=g-importing-member from mujoco.mjx._src import math from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import Model from mujoco.mjx._src.types import ObjType from mujoco.mjx._src.types import SensorType @@ -30,6 +31,9 @@ import numpy as np def sensor_pos(m: Model, d: Data) -> Data: """Compute position-dependent sensors values.""" + if m.opt.disableflags & DisableBit.SENSOR: + return d + # no position-dependent sensors stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS if sum(stage_pos) == 0: @@ -145,9 +149,17 @@ def sensor_pos(m: Model, d: Data) -> Data: def sensor_vel(m: Model, d: Data) -> Data: """Compute velocity-dependent sensors values.""" + + if m.opt.disableflags & DisableBit.SENSOR: + return d + return d def sensor_acc(m: Model, d: Data) -> Data: """Compute acceleration/force-dependent sensors values.""" + + if m.opt.disableflags & DisableBit.SENSOR: + return d + return d diff --git a/mjx/mujoco/mjx/_src/sensor_test.py b/mjx/mujoco/mjx/_src/sensor_test.py index d120e757..0ff11a86 100644 --- a/mjx/mujoco/mjx/_src/sensor_test.py +++ b/mjx/mujoco/mjx/_src/sensor_test.py @@ -17,6 +17,7 @@ from absl.testing import absltest from absl.testing import parameterized import jax +from jax import numpy as jp import mujoco from mujoco import mjx from mujoco.mjx._src import test_util @@ -62,6 +63,25 @@ class SensorTest(parameterized.TestCase): # sensor values _assert_eq(d.sensordata, dx.sensordata, 'sensordata') + def test_disable_sensor(self): + """Tests disabling sensor.""" + m = test_util.load_test_file('sensor.xml') + # disable sensors + m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.SENSOR + d = mujoco.MjData(m) + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) + mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) + # random sensor values + random_sensor = jp.array(np.random.random(dx.sensordata.shape)) + dx = dx.replace(sensordata=random_sensor) + # call sensor functions + dx = jax.jit(mjx.forward)(mx, dx) + # sensor values + _assert_eq(random_sensor, dx.sensordata, 'sensordata') + def test_unsupported_sensor(self): """Tests MJX sensor functions do not break for unsupported sensors.""" m = test_util.load_test_file('unsupported_sensor.xml') diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index c3f2d2ff..8ac87e41 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -45,6 +45,7 @@ class DisableBit(enum.IntFlag): WARMSTART: warmstart constraint solver ACTUATION: apply actuation forces REFSAFE: integrator safety: make ref[0]>=2*timestep + SENSOR: sensors """ CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT EQUALITY = mujoco.mjtDisableBit.mjDSBL_EQUALITY @@ -56,9 +57,10 @@ class DisableBit(enum.IntFlag): WARMSTART = mujoco.mjtDisableBit.mjDSBL_WARMSTART ACTUATION = mujoco.mjtDisableBit.mjDSBL_ACTUATION REFSAFE = mujoco.mjtDisableBit.mjDSBL_REFSAFE + SENSOR = mujoco.mjtDisableBit.mjDSBL_SENSOR EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP FILTERPARENT = mujoco.mjtDisableBit.mjDSBL_FILTERPARENT - # unsupported: FRICTIONLOSS, SENSOR, MIDPHASE + # unsupported: FRICTIONLOSS, MIDPHASE class JointType(enum.IntEnum): From 51733c2a8ac081ed7c44d8e96bb0f8d6d30059f0 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Sun, 18 Aug 2024 05:15:10 -0700 Subject: [PATCH 03/26] Add rangefinder sensor to MJX. PiperOrigin-RevId: 664417458 Change-Id: I1ad7720e48541349f155984584314b2b49ccbf36 --- doc/changelog.rst | 4 ++-- mjx/mujoco/mjx/_src/sensor.py | 14 ++++++++++++++ mjx/mujoco/mjx/_src/types.py | 2 ++ mjx/mujoco/mjx/test_data/sensor.xml | 15 ++++++++++++++- 4 files changed, 32 insertions(+), 3 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 1fd5ebf1..9152b8a7 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -17,8 +17,8 @@ General MJX ^^^ 5. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`). -6. Added position-dependent sensors: ``MAGNETOMETER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, - ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``. +6. Added position-dependent sensors: ``MAGNETOMETER``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, + ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``. 7. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. 8. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index e644bcdc..afe3741d 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -19,6 +19,7 @@ from jax import numpy as jp import mujoco # pylint: disable=g-importing-member from mujoco.mjx._src import math +from mujoco.mjx._src import ray from mujoco.mjx._src.types import Data from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import Model @@ -71,6 +72,19 @@ def sensor_pos(m: Model, d: Data) -> Data: d.site_xmat[objid] ).reshape(-1) adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) + elif sensor_type == SensorType.RANGEFINDER: + site_bodyid = m.site_bodyid[objid] + for sid in set(site_bodyid): + id_ = sid == site_bodyid + objid_ = objid[id_] + site_xpos = d.site_xpos[objid_] + site_mat = d.site_xmat[objid_].reshape((-1, 9))[:, np.array([2, 5, 8])] + sensor, _ = jax.vmap( + ray.ray, in_axes=(None, None, 0, 0, None, None, None) + )(m, d, site_xpos, site_mat, (), True, sid) + sensors.append(sensor) + adrs.append(adr[id_]) + continue # avoid adding to sensors/adrs list a second time elif sensor_type == SensorType.JOINTPOS: sensor = d.qpos[m.jnt_qposadr[objid]] elif sensor_type == SensorType.ACTUATORPOS: diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 8ac87e41..b92db752 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -293,6 +293,7 @@ class SensorType(enum.IntEnum): Members: MAGNETOMETER: magnetometer + RANGEFINDER: rangefinder JOINTPOS: joint position ACTUATORPOS: actuator position BALLQUAT: ball joint orientation @@ -304,6 +305,7 @@ class SensorType(enum.IntEnum): CLOCK: simulation time """ MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER + RANGEFINDER = mujoco.mjtSensor.mjSENS_RANGEFINDER JOINTPOS = mujoco.mjtSensor.mjSENS_JOINTPOS ACTUATORPOS = mujoco.mjtSensor.mjSENS_ACTUATORPOS BALLQUAT = mujoco.mjtSensor.mjSENS_BALLQUAT diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml index 13178604..80763173 100644 --- a/mjx/mujoco/mjx/test_data/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor.xml @@ -2,6 +2,7 @@ * position-dependent sensors: -magnetometer +-rangefinder -jointpos -actuatorpos -ballquat @@ -15,11 +16,16 @@ * acceleration/force-dependent sensors: --> + + + - + + + @@ -39,6 +45,11 @@ + + + + + @@ -51,6 +62,7 @@ + @@ -67,6 +79,7 @@ + From 6a12787a3200a698d0e93a7d0e8e228ad96ce12f Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Sun, 18 Aug 2024 05:56:00 -0700 Subject: [PATCH 04/26] Add camera projection sensor to MJX. PiperOrigin-RevId: 664426007 Change-Id: I0f2681c156b12ab62503d45beb2cd8752e7202d6 --- doc/changelog.rst | 5 ++- mjx/mujoco/mjx/_src/sensor.py | 60 +++++++++++++++++++++++++++++ mjx/mujoco/mjx/_src/types.py | 2 + mjx/mujoco/mjx/test_data/sensor.xml | 19 +++++++-- 4 files changed, 81 insertions(+), 5 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 9152b8a7..4ddf86f1 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -17,8 +17,9 @@ General MJX ^^^ 5. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`). -6. Added position-dependent sensors: ``MAGNETOMETER``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, - ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``. +6. Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, + ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, + ``CLOCK``. 7. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. 8. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index afe3741d..f8792764 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -72,6 +72,66 @@ def sensor_pos(m: Model, d: Data) -> Data: d.site_xmat[objid] ).reshape(-1) adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) + elif sensor_type == SensorType.CAMPROJECTION: + + @jax.vmap + def _cam_project( + target_xpos, xpos, xmat, res, fovy, intrinsic, sensorsize, focal_flag + ): + translation = jp.eye(4).at[0:3, 3].set(-xpos) + rotation = jp.eye(4).at[:3, :3].set(xmat.T) + + # focal transformation matrix (3 x 4) + f = 0.5 / jp.tan(fovy * jp.pi / 360.0) * res[1] + fx, fy = jp.where( + focal_flag, + intrinsic[:2] / (sensorsize[:2] + mujoco.mjMINVAL) * res[:2], + f, + ) # add mjMINVAL to denominator to prevent divide by zero warning + + focal = jp.array([[-fx, 0, 0, 0], [0, fy, 0, 0], [0, 0, 1.0, 0]]) + + # image matrix (3 x 3) + image = jp.eye(3).at[:2, 2].set(res[0:2] / 2.0) + + # projection matrix (3 x 4): product of all 4 matrices + proj = image @ focal @ rotation @ translation + + # projection matrix multiplies homogenous [x, y, z, 1] vectors + pos_hom = jp.append(target_xpos, 1.0) + + # project world coordinates into pixel space, see: + # https://en.wikipedia.org/wiki/3D_projection#Mathematical_formula + pixel_coord_hom = proj @ pos_hom + + # avoid dividing by tiny numbers + denom = pixel_coord_hom[2] + denom = jp.where( + jp.abs(denom) < mujoco.mjMINVAL, + jp.clip(denom, -mujoco.mjMINVAL, mujoco.mjMINVAL), + denom, + ) + + # compute projection + sensor = pixel_coord_hom / denom + + return sensor[:2] + + refid = m.sensor_refid[idx] + sensorsize = m.cam_sensorsize[refid] + intrinsic = m.cam_intrinsic[refid] + fovy = m.cam_fovy[refid] + res = m.cam_resolution[refid] + focal_flag = np.logical_and(sensorsize[:, 0] != 0, sensorsize[:, 1] != 0) + + target_xpos = d.site_xpos[objid] + xpos = d.cam_xpos[refid] + xmat = d.cam_xmat[refid] + + sensor = _cam_project( + target_xpos, xpos, xmat, res, fovy, intrinsic, sensorsize, focal_flag + ).reshape(-1) + adr = (adr[:, None] + np.arange(2)[None]).reshape(-1) elif sensor_type == SensorType.RANGEFINDER: site_bodyid = m.site_bodyid[objid] for sid in set(site_bodyid): diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index b92db752..de53a7a1 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -293,6 +293,7 @@ class SensorType(enum.IntEnum): Members: MAGNETOMETER: magnetometer + CAMPROJECTION: camera projection RANGEFINDER: rangefinder JOINTPOS: joint position ACTUATORPOS: actuator position @@ -305,6 +306,7 @@ class SensorType(enum.IntEnum): CLOCK: simulation time """ MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER + CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION RANGEFINDER = mujoco.mjtSensor.mjSENS_RANGEFINDER JOINTPOS = mujoco.mjtSensor.mjSENS_JOINTPOS ACTUATORPOS = mujoco.mjtSensor.mjSENS_ACTUATORPOS diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml index 80763173..72b21e3f 100644 --- a/mjx/mujoco/mjx/test_data/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor.xml @@ -2,6 +2,7 @@ * position-dependent sensors: -magnetometer +-camprojection -rangefinder -jointpos -actuatorpos @@ -46,6 +47,19 @@ + + + + + + + + + + + + @@ -60,7 +74,6 @@ - @@ -71,6 +84,7 @@ + @@ -78,11 +92,10 @@ + - - From 6241814ad79d766c6ff51067ba000f1196d09a08 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Sun, 18 Aug 2024 08:46:45 -0700 Subject: [PATCH 05/26] Add velocity-dependent sensors to MJX. Includes: joint velocity, actuator velocity, and ball joint angular velocity. PiperOrigin-RevId: 664461284 Change-Id: I5063f662d707fbe342493f9cd7993f009fac9c76 --- doc/changelog.rst | 11 ++++---- mjx/mujoco/mjx/_src/sensor.py | 39 ++++++++++++++++++++++++----- mjx/mujoco/mjx/_src/types.py | 6 +++++ mjx/mujoco/mjx/test_data/sensor.xml | 9 +++++++ 4 files changed, 54 insertions(+), 11 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 4ddf86f1..87cbebea 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -20,19 +20,20 @@ MJX 6. Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``. -7. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. -8. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. +7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. +8. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. +9. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. Bug fixes ^^^^^^^^^ -9. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, +10. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, contribution by :github:user:`michael-ahn`). -10. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit +11. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit integrators, wrong derivatives would be computed. Python bindings ^^^^^^^^^^^^^^^ -11. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). +12. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). Version 3.2.2 (Aug 8, 2024) diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index f8792764..77f6a3aa 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -35,11 +35,6 @@ def sensor_pos(m: Model, d: Data) -> Data: if m.opt.disableflags & DisableBit.SENSOR: return d - # no position-dependent sensors - stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS - if sum(stage_pos) == 0: - return d - # position and orientation by object type objtype_data = { ObjType.UNKNOWN: ( @@ -60,6 +55,7 @@ def sensor_pos(m: Model, d: Data) -> Data: SensorType.FRAMEZAXIS: 2, } + stage_pos = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS sensors, adrs = [], [] for sensor_type in set(m.sensor_type[stage_pos]): @@ -206,6 +202,7 @@ def sensor_pos(m: Model, d: Data) -> Data: elif sensor_type == SensorType.CLOCK: sensor = jp.repeat(d.time, sum(idx)) else: + # TODO(taylorhowell): raise error after adding sensor check to io.py continue # unsupported sensor type sensors.append(sensor) @@ -227,7 +224,37 @@ def sensor_vel(m: Model, d: Data) -> Data: if m.opt.disableflags & DisableBit.SENSOR: return d - return d + stage_vel = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_VEL + sensors, adrs = [], [] + + for sensor_type in set(m.sensor_type[stage_vel]): + idx = m.sensor_type == sensor_type + objid = m.sensor_objid[idx] + adr = m.sensor_adr[idx] + + if sensor_type == SensorType.JOINTVEL: + sensor = d.qvel[m.jnt_dofadr[objid]] + elif sensor_type == SensorType.ACTUATORVEL: + sensor = d.actuator_velocity[objid] + elif sensor_type == SensorType.BALLANGVEL: + jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None] + sensor = d.qvel[jnt_dotadr].reshape(-1) + adr = (adr[:, None] + np.arange(3)[None]).reshape(-1) + else: + # TODO(taylorhowell): raise error after adding sensor check to io.py + continue # unsupported sensor typ + + sensors.append(sensor) + adrs.append(adr) + + if not adrs: + return d + + sensordata = d.sensordata.at[np.concatenate(adrs)].set( + jp.concatenate(sensors) + ) + + return d.replace(sensordata=sensordata) def sensor_acc(m: Model, d: Data) -> Data: diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index de53a7a1..15e8bb18 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -304,6 +304,9 @@ class SensorType(enum.IntEnum): FRAMEZAXIS: frame z-axis SUBTREECOM: subtree centor of mass CLOCK: simulation time + JOINTVEL: joint velocity + ACTUATORVEL: actuator velocity + BALLANGVEL: ball joint angular velocity """ MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION @@ -317,6 +320,9 @@ class SensorType(enum.IntEnum): FRAMEZAXIS = mujoco.mjtSensor.mjSENS_FRAMEZAXIS SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM CLOCK = mujoco.mjtSensor.mjSENS_CLOCK + JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL + ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL + BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL class ObjType(PyTreeNode): diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml index 72b21e3f..04924952 100644 --- a/mjx/mujoco/mjx/test_data/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor.xml @@ -14,6 +14,9 @@ -subtreecom -clock * velocity-dependent sensors: +-jointvel +-actuatorvel +-ballangvel * acceleration/force-dependent sensors: --> @@ -77,8 +80,11 @@ + + + @@ -87,12 +93,15 @@ + + + From 0fa39164f744244f2cc1ac584216791a4bb996e1 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 19 Aug 2024 07:54:01 -0700 Subject: [PATCH 06/26] Speed up mesh support function with warmstarts. PiperOrigin-RevId: 664807396 Change-Id: I9124e65efd9140a8e858b7b6be0867c26961477b --- src/engine/engine_collision_convex.c | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/engine/engine_collision_convex.c b/src/engine/engine_collision_convex.c index e74b29d3..ceba996b 100644 --- a/src/engine/engine_collision_convex.c +++ b/src/engine/engine_collision_convex.c @@ -219,11 +219,11 @@ void mjc_support(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) { vert_globalid = m->mesh_graph + graphadr + 2 + numvert; edge_localid = m->mesh_graph + graphadr + 2 + 2*numvert; - // init with first vertex in convex hull - ibest = 0; - tmp = local_dir[0] * (mjtNum)vertdata[3*vert_globalid[0]] + - local_dir[1] * (mjtNum)vertdata[3*vert_globalid[0]+1] + - local_dir[2] * (mjtNum)vertdata[3*vert_globalid[0]+2]; + // init with first vertex in convex hull or warmstart + ibest = obj->meshindex < 0 ? 0 : obj->meshindex; + tmp = local_dir[0] * (mjtNum)vertdata[3*vert_globalid[ibest]+0] + + local_dir[1] * (mjtNum)vertdata[3*vert_globalid[ibest]+1] + + local_dir[2] * (mjtNum)vertdata[3*vert_globalid[ibest]+2]; // hill-climb until no change change = 1; From 4aab00fa2df3ba56ac2dd4e23d076b4e3027cc9b Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 19 Aug 2024 08:42:05 -0700 Subject: [PATCH 07/26] Minor cleanups to body and joint compilers. PiperOrigin-RevId: 664824939 Change-Id: I6014ff10683849e2926f1b1c8df8b6862192cd00 --- src/user/user_objects.cc | 59 ++++++++++++++++++---------------------- src/user/user_objects.h | 2 +- 2 files changed, 27 insertions(+), 34 deletions(-) diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index cc0c7c2f..77a9b896 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -1307,7 +1307,7 @@ mjsElement* mjCBody::NextChild(mjsElement* child, mjtObj type) { // compute geom inertial frame: ipos, iquat, mass, inertia -void mjCBody::GeomFrame(void) { +void mjCBody::InertiaFromGeom(void) { int sz; double com[3] = {0, 0, 0}; double toti[6] = {0, 0, 0, 0, 0, 0}; @@ -1477,7 +1477,7 @@ void mjCBody::Compile(void) { throw mjCError(this, "error '%s' in inertia alternative", ierr); } - // compile all geoms, phase 1 + // compile all geoms for (int i=0; iinferinertia = id>0 && (!explicitinertial || model->inertiafromgeom == mjINERTIAFROMGEOM_TRUE) && @@ -1489,7 +1489,7 @@ void mjCBody::Compile(void) { // set inertial frame from geoms if necessary if (id>0 && (model->inertiafromgeom==mjINERTIAFROMGEOM_TRUE || (!mjuu_defined(ipos[0]) && model->inertiafromgeom==mjINERTIAFROMGEOM_AUTO))) { - GeomFrame(); + InertiaFromGeom(); } // both pos and ipos undefined: error @@ -1574,17 +1574,16 @@ void mjCBody::Compile(void) { } // make sure mocap body is fixed child of world - if (mocap) - if (dofnum || parentid) { - throw mjCError(this, "mocap body '%s' is not a fixed child of world", name.c_str()); - } + if (mocap && (dofnum || parentid)) { + throw mjCError(this, "mocap body '%s' is not a fixed child of world", name.c_str()); + } // compute body global pose (no joint transformations in qpos0) if (id>0) { - mjCBody* par = model->Bodies()[parentid]; - mjuu_rotVecQuat(xpos0, pos, par->xquat0); - mjuu_addtovec(xpos0, par->xpos0, 3); - mjuu_mulquat(xquat0, par->xquat0, quat); + mjCBody* parent = model->Bodies()[parentid]; + mjuu_rotVecQuat(xpos0, pos, parent->xquat0); + mjuu_addtovec(xpos0, parent->xpos0, 3); + mjuu_mulquat(xquat0, parent->xquat0, quat); } // compile all sites @@ -1613,15 +1612,13 @@ void mjCBody::Compile(void) { } } - if (!model->discardvisual) { - return; - } - - // set inertial to explicit for bodies containing visual geoms - for (int j=0; jIsVisual()) { - explicitinertial = true; - break; + // if discarding visual geoms, use explicit inertias + if (model->discardvisual) { + for (int j=0; jIsVisual()) { + explicitinertial = true; + break; + } } } } @@ -1946,22 +1943,15 @@ int mjCJoint::Compile(void) { } } - // frame - if (frame) { - double mat[9]; - mjuu_quat2mat(mat, frame->quat); - mjuu_mulvecmat(axis, axis, mat); - } - - // FREE or BALL: set axis to (0,0,1) + // axis: FREE or BALL are fixed to (0,0,1) if (type==mjJNT_FREE || type==mjJNT_BALL) { axis[0] = axis[1] = 0; axis[2] = 1; } - // FREE: set pos to (0,0,0) - if (type==mjJNT_FREE) { - mjuu_zerovec(pos, 3); + // otherwise accumulate frame rotation + else if (frame) { + mjuu_rotVecQuat(axis, axis, frame->quat); } // normalize axis, check norm @@ -1974,10 +1964,13 @@ int mjCJoint::Compile(void) { throw mjCError(this, "limits should not be defined in free joint"); } - // compute local position + // pos: FREE is fixed to (0,0,0) if (type == mjJNT_FREE) { mjuu_zerovec(pos, 3); - } else if (frame) { + } + + // otherwise accumulate frame translation + else if (frame) { double qunit[4] = {1, 0, 0, 0}; mjuu_frameaccumChild(frame->pos, frame->quat, pos, qunit); } diff --git a/src/user/user_objects.h b/src/user/user_objects.h index c55cdcd0..47442f1d 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -336,7 +336,7 @@ class mjCBody : public mjCBody_, private mjsBody { mjCBody& operator=(const mjCBody& other); // copy assignment void Compile(void); // compiler - void GeomFrame(void); // get inertial info from geoms + void InertiaFromGeom(void); // get inertial info from geoms // objects allocated by Add functions std::vector bodies; // child bodies From 5d91231d3195113190aa89ebe7d1e1f4b175736a Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 19 Aug 2024 10:37:50 -0700 Subject: [PATCH 08/26] Rename `mjBuffer` to `mjByteVec`. PiperOrigin-RevId: 664872558 Change-Id: I7b003374c390958667635fb21089faf2ece54b14 --- doc/APIreference/APItypes.rst | 6 +++--- doc/includes/references.h | 4 ++-- include/mujoco/mjspec.h | 6 +++--- include/mujoco/mujoco.h | 2 +- introspect/functions.py | 2 +- introspect/structs.py | 2 +- python/mujoco/codegen/generate_spec_bindings.py | 2 +- src/user/user_api.cc | 2 +- src/user/user_api.h | 2 +- 9 files changed, 14 insertions(+), 14 deletions(-) diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst index 9be6d9e4..5f58d803 100644 --- a/doc/APIreference/APItypes.rst +++ b/doc/APIreference/APItypes.rst @@ -1393,7 +1393,7 @@ Alternative orientation specifiers. .. _ArrayHandles: -.. _mjBuffer: +.. _mjByteVec: .. _mjString: @@ -1419,7 +1419,6 @@ C handles for C++ strings and vector types. When using from C, use the provided #ifdef __cplusplus // C++: defined to be compatible with corresponding std types - using mjBuffer = std::vector; using mjString = std::string; using mjStringVec = std::vector; using mjIntVec = std::vector; @@ -1427,9 +1426,9 @@ C handles for C++ strings and vector types. When using from C, use the provided using mjFloatVec = std::vector; using mjFloatVecVec = std::vector>; using mjDoubleVec = std::vector; + using mjByteVec = std::vector; #else // C: opaque types - typedef void mjBuffer; typedef void mjString; typedef void mjStringVec; typedef void mjIntVec; @@ -1437,6 +1436,7 @@ C handles for C++ strings and vector types. When using from C, use the provided typedef void mjFloatVec; typedef void mjFloatVecVec; typedef void mjDoubleVec; + typedef void mjByteVec; #endif diff --git a/doc/includes/references.h b/doc/includes/references.h index f0bc5390..01942d4e 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -2028,7 +2028,7 @@ typedef struct mjsTexture_ { // texture specification mjStringVec* cubefiles; // different file for each side of the cube // method 4: from buffer read by user - mjBuffer* data; // texture data + mjByteVec* data; // texture data // flip options mjtByte hflip; // horizontal flip @@ -3588,7 +3588,7 @@ mjsHField* mjs_asHField(mjsElement* element); mjsSkin* mjs_asSkin(mjsElement* element); mjsTexture* mjs_asTexture(mjsElement* element); mjsMaterial* mjs_asMaterial(mjsElement* element); -void mjs_setBuffer(mjBuffer* dest, const void* array, int size); +void mjs_setBuffer(mjByteVec* dest, const void* array, int size); void mjs_setString(mjString* dest, const char* text); void mjs_setStringVec(mjStringVec* dest, const char* text); mjtByte mjs_setInStringVec(mjStringVec* dest, int i, const char* text); diff --git a/include/mujoco/mjspec.h b/include/mujoco/mjspec.h index 1057e559..7d76a46e 100644 --- a/include/mujoco/mjspec.h +++ b/include/mujoco/mjspec.h @@ -33,7 +33,6 @@ extern "C" { #ifdef __cplusplus // C++: defined to be compatible with corresponding std types - using mjBuffer = std::vector; using mjString = std::string; using mjStringVec = std::vector; using mjIntVec = std::vector; @@ -41,9 +40,9 @@ extern "C" { using mjFloatVec = std::vector; using mjFloatVecVec = std::vector>; using mjDoubleVec = std::vector; + using mjByteVec = std::vector; #else // C: opaque types - typedef void mjBuffer; typedef void mjString; typedef void mjStringVec; typedef void mjIntVec; @@ -51,6 +50,7 @@ extern "C" { typedef void mjFloatVec; typedef void mjFloatVecVec; typedef void mjDoubleVec; + typedef void mjByteVec; #endif @@ -515,7 +515,7 @@ typedef struct mjsTexture_ { // texture specification mjStringVec* cubefiles; // different file for each side of the cube // method 4: from buffer read by user - mjBuffer* data; // texture data + mjByteVec* data; // texture data // flip options mjtByte hflip; // horizontal flip diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index ac8a5507..86897462 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1630,7 +1630,7 @@ MJAPI mjsMaterial* mjs_asMaterial(mjsElement* element); //---------------------------------- Attribute setters --------------------------------------------- // Copy buffer. -MJAPI void mjs_setBuffer(mjBuffer* dest, const void* array, int size); +MJAPI void mjs_setBuffer(mjByteVec* dest, const void* array, int size); // Copy text to string. MJAPI void mjs_setString(mjString* dest, const char* text); diff --git a/introspect/functions.py b/introspect/functions.py index bf2860ba..4e3073e8 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -10253,7 +10253,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ FunctionParameterDecl( name='dest', type=PointerType( - inner_type=ValueType(name='mjBuffer'), + inner_type=ValueType(name='mjByteVec'), ), ), FunctionParameterDecl( diff --git a/introspect/structs.py b/introspect/structs.py index 12d2af2b..2ed6deef 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -10014,7 +10014,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ StructFieldDecl( name='data', type=PointerType( - inner_type=ValueType(name='mjBuffer'), + inner_type=ValueType(name='mjByteVec'), ), doc='texture data', ), diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py index 50b40d1b..2ac03612 100644 --- a/python/mujoco/codegen/generate_spec_bindings.py +++ b/python/mujoco/codegen/generate_spec_bindings.py @@ -148,7 +148,7 @@ def _ptr_binding_code( self.{fullvarname}->push_back(py::cast<{vartype}>(val)); }} }}, py::return_value_policy::reference_internal);""" - elif vartype == 'mjBuffer': # C++ buffer -> Python list + elif vartype == 'mjByteVec': # C++ buffer -> Python list return f"""\ {classname}.def_property( "{varname}", diff --git a/src/user/user_api.cc b/src/user/user_api.cc index 4743f83b..f598db61 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -885,7 +885,7 @@ mjsMaterial* mjs_asMaterial(mjsElement* element) { // copy buffer to destination buffer -void mjs_setBuffer(mjBuffer* dest, const void* array, int size) { +void mjs_setBuffer(mjByteVec* dest, const void* array, int size) { const std::byte* buffer = static_cast(array); dest->clear(); dest->reserve(size); diff --git a/src/user/user_api.h b/src/user/user_api.h index 967320f5..c30abd53 100644 --- a/src/user/user_api.h +++ b/src/user/user_api.h @@ -303,7 +303,7 @@ MJAPI mjsMaterial* mjs_asMaterial(mjsElement* element); //---------------------------------- Attribute setters --------------------------------------------- // Copy buffer. -MJAPI void mjs_setBuffer(mjBuffer* dest, const void* array, int size); +MJAPI void mjs_setBuffer(mjByteVec* dest, const void* array, int size); // Copy text to string. MJAPI void mjs_setString(mjString* dest, const char* text); From a68141eeffaac14f69e8464ab1ca477bc90e0cc5 Mon Sep 17 00:00:00 2001 From: Erik Frey Date: Mon, 19 Aug 2024 11:11:34 -0700 Subject: [PATCH 09/26] Implicitfast integration for MJX. PiperOrigin-RevId: 664889120 Change-Id: Id3bd3916fe821ab2af79c52e53e14837ae829b2a --- doc/changelog.rst | 8 +++-- doc/mjx.rst | 4 +-- mjx/mujoco/mjx/__init__.py | 3 +- mjx/mujoco/mjx/_src/forward.py | 42 ++++++++++++++++++++++++ mjx/mujoco/mjx/_src/forward_test.py | 6 ++++ mjx/mujoco/mjx/_src/io.py | 7 +++- mjx/mujoco/mjx/_src/io_test.py | 8 +++++ mjx/mujoco/mjx/_src/types.py | 4 ++- mjx/mujoco/mjx/test_data/constraints.xml | 2 +- 9 files changed, 75 insertions(+), 9 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 87cbebea..9b47e771 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -23,17 +23,19 @@ MJX 7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. 8. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. 9. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. +10. Added support for :ref:`implicitfast integration` for all cases except + :doc:`fluid drag `. Bug fixes ^^^^^^^^^ -10. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, +11. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, contribution by :github:user:`michael-ahn`). -11. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit +12. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit integrators, wrong derivatives would be computed. Python bindings ^^^^^^^^^^^^^^^ -12. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). +13. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). Version 3.2.2 (Aug 8, 2024) diff --git a/doc/mjx.rst b/doc/mjx.rst index 9aafa764..90a43e87 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -202,7 +202,7 @@ The following features are **fully supported** in MJX: * - :ref:`Equality ` - ``CONNECT``, ``WELD``, ``JOINT``, ``TENDON`` * - :ref:`Integrator ` - - ``EULER``, ``RK4`` + - ``EULER``, ``RK4``, ``IMPLICITFAST`` (``IMPLICITFAST`` not supported with :doc:`fluid drag `) * - :ref:`Cone ` - ``PYRAMIDAL``, ``ELLIPTIC`` * - :ref:`Condim ` @@ -229,7 +229,7 @@ The following features are **in development** and coming soon: * - :ref:`Constraint ` - :ref:`Frictionloss `, ``FRICTION_DOF`` * - :ref:`Integrator ` - - ``IMPLICIT``, ``IMPLICITFAST`` + - ``IMPLICIT`` * - Dynamics - :ref:`Inverse ` * - :ref:`Actuator Dynamics ` diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 2b442d77..a4a8a5da 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -23,6 +23,7 @@ from mujoco.mjx._src.forward import fwd_acceleration from mujoco.mjx._src.forward import fwd_actuation from mujoco.mjx._src.forward import fwd_position from mujoco.mjx._src.forward import fwd_velocity +from mujoco.mjx._src.forward import implicit from mujoco.mjx._src.forward import rungekutta4 from mujoco.mjx._src.forward import step from mujoco.mjx._src.io import get_data @@ -32,9 +33,9 @@ from mujoco.mjx._src.io import put_data from mujoco.mjx._src.io import put_model from mujoco.mjx._src.passive import passive from mujoco.mjx._src.ray import ray +from mujoco.mjx._src.sensor import sensor_acc from mujoco.mjx._src.sensor import sensor_pos from mujoco.mjx._src.sensor import sensor_vel -from mujoco.mjx._src.sensor import sensor_acc from mujoco.mjx._src.smooth import camlight from mujoco.mjx._src.smooth import com_pos from mujoco.mjx._src.smooth import com_vel diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index a90ef854..456e748c 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -348,6 +348,46 @@ def rungekutta4(m: Model, d: Data) -> Data: return d +@named_scope +def implicit(m: Model, d: Data) -> Data: + """Integrates fully implicit in velocity.""" + + qderiv = None + + # qDeriv += d qfrc_actuator / d qvel + if not m.opt.disableflags & DisableBit.ACTUATION: + affine_bias = m.actuator_biastype == BiasType.AFFINE + bias_vel = m.actuator_biasprm[:, 2] * affine_bias + affine_gain = m.actuator_gaintype == GainType.AFFINE + gain_vel = m.actuator_gainprm[:, 2] * affine_gain + ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act) + vel = bias_vel + gain_vel * ctrl + qderiv = d.actuator_moment.T @ jp.diag(vel) @ d.actuator_moment + + # qDeriv += d qfrc_passive / d qvel + if not m.opt.disableflags & DisableBit.PASSIVE: + if qderiv is None: + qderiv = -jp.diag(m.dof_damping) + else: + qderiv -= jp.diag(m.dof_damping) + if m.ntendon: + qderiv -= d.ten_J.T @ jp.diag(m.tendon_damping) @ d.ten_J + # TODO(robotics-simulation): fluid drag model + if m.opt.has_fluid_params: + raise NotImplementedError('fluid drag not supported for implicitfast') + + qacc = d.qacc + if qderiv is not None: + # TODO(robotics-simulation): use smooth.factor_m / solve_m here: + qm = support.full_m(m, d) if support.is_sparse(m) else d.qM + qm -= m.opt.timestep * qderiv + qh, _ = jax.scipy.linalg.cho_factor(qm) + qfrc = d.qfrc_smooth + d.qfrc_constraint + qacc = jax.scipy.linalg.cho_solve((qh, False), qfrc) + + return _advance(m, d, d.act_dot, qacc) + + @named_scope def forward(m: Model, d: Data) -> Data: """Forward dynamics.""" @@ -377,6 +417,8 @@ def step(m: Model, d: Data) -> Data: d = euler(m, d) elif m.opt.integrator == IntegratorType.RK4: d = rungekutta4(m, d) + elif m.opt.integrator == IntegratorType.IMPLICITFAST: + d = implicit(m, d) else: raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.') diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py index 0a777797..a93f04e0 100644 --- a/mjx/mujoco/mjx/_src/forward_test.py +++ b/mjx/mujoco/mjx/_src/forward_test.py @@ -68,6 +68,12 @@ class ForwardTest(absltest.TestCase): _assert_attr_eq(d, dx, 'qpos') _assert_attr_eq(d, dx, 'time') + # implicitfast + m.opt.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST + dx = jax.jit(mjx.implicit)(mx, mjx.put_data(m, d)) + mujoco.mj_implicit(m, d) + _assert_attr_eq(d, dx, 'qpos') + def test_step(self): m = test_util.load_test_file('constraints.xml') d = mujoco.MjData(m) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 7ad1c261..4caf0be9 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -47,13 +47,18 @@ def _make_option(o: mujoco.MjOption) -> types.Option: if o.enableflags & 2**i: raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}') + has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any() + implicitfast = o.integrator == mujoco.mjtIntegrator.mjINT_IMPLICITFAST + if implicitfast and has_fluid_params: + raise NotImplementedError('implicitfast not implemented for fluid drag.') + fields = {f.name: getattr(o, f.name, None) for f in types.Option.fields()} fields['integrator'] = types.IntegratorType(o.integrator) fields['cone'] = types.ConeType(o.cone) fields['jacobian'] = types.JacobianType(o.jacobian) fields['solver'] = types.SolverType(o.solver) fields['disableflags'] = types.DisableBit(o.disableflags) - fields['has_fluid_params'] = o.density > 0 or o.viscosity > 0 or o.wind.any() + fields['has_fluid_params'] = has_fluid_params return types.Option(**fields) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 020e8e23..97c61382 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -208,6 +208,14 @@ class ModelIOTest(parameterized.TestCase): """)) + def test_implicitfast_fluid_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + """)) + class DataIOTest(parameterized.TestCase): """IO tests for mjx.Data.""" diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 15e8bb18..4a14db17 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -90,10 +90,12 @@ class IntegratorType(enum.IntEnum): Members: EULER: semi-implicit Euler RK4: 4th-order Runge Kutta + IMPLICITFAST: implicit in velocity, no rne derivative """ EULER = mujoco.mjtIntegrator.mjINT_EULER RK4 = mujoco.mjtIntegrator.mjINT_RK4 - # unsupported: IMPLICIT, IMPLICITFAST + IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST + # unsupported: IMPLICIT class GeomType(enum.IntEnum): diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml index 8f7e5643..a1476a45 100644 --- a/mjx/mujoco/mjx/test_data/constraints.xml +++ b/mjx/mujoco/mjx/test_data/constraints.xml @@ -83,7 +83,7 @@ - + From 029fbd5172bdc37df6f2abe2aafa36e515dd6ecc Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 20 Aug 2024 12:13:10 +0000 Subject: [PATCH 10/26] Bump zipp from 3.17.0 to 3.19.1 in /mjx Bumps [zipp](https://github.com/jaraco/zipp) from 3.17.0 to 3.19.1. - [Release notes](https://github.com/jaraco/zipp/releases) - [Changelog](https://github.com/jaraco/zipp/blob/main/NEWS.rst) - [Commits](https://github.com/jaraco/zipp/compare/v3.17.0...v3.19.1) --- updated-dependencies: - dependency-name: zipp dependency-type: direct:production ... Signed-off-by: dependabot[bot] --- mjx/requirements.txt | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mjx/requirements.txt b/mjx/requirements.txt index f3aadfb8..fe5c2623 100644 --- a/mjx/requirements.txt +++ b/mjx/requirements.txt @@ -83,8 +83,9 @@ importlib-resources==6.1.0 \ --hash=sha256:aa50258bbfa56d4e33fbd8aa3ef48ded10d1735f11532b8df95388cc6bdb7e83 typing_extensions==4.8.0 \ --hash=sha256:8f92fc8806f9a6b641eaa5318da32b44d401efaac0f6678c9bc448ba3605faa0 -zipp==3.17.0 \ - --hash=sha256:0e923e726174922dce09c53c59ad483ff7bbb8e572e00c7f7c46b88556409f31 +zipp==3.19.1 \ + --hash=sha256:2828e64edb5386ea6a52e7ba7cdb17bb30a73a858f5eb6eb93d8d36f5ea26091 \ + --hash=sha256:35427f6d5594f4acf82d25541438348c26736fa9b3afa2754bcd63cdb99d8e8f # Transitive dependencies of jax and jaxlib importlib-metadata==6.8.0; python_version < '3.10' \ From ad9cd7301e06c2a3aef4c902acaa149de5af04b3 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Tue, 20 Aug 2024 06:47:17 -0700 Subject: [PATCH 11/26] Correctly apply frame transform to light direction. PiperOrigin-RevId: 665340267 Change-Id: Ida16ddbf3a057ffd3e260b71180a45acbf749931 --- src/user/user_objects.cc | 9 ++++++--- test/user/user_objects_test.cc | 29 ++++++++++++++++++++++++++++- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 77a9b896..2fe975b4 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -3222,11 +3222,14 @@ void mjCLight::CopyFromSpec() { void mjCLight::Compile(void) { CopyFromSpec(); - double quat[4]= {1, 0, 0, 0}; - // frame if (frame) { - mjuu_frameaccumChild(frame->pos, frame->quat, pos, quat); + // apply frame transform to pos, qunit is unused + double qunit[4]= {1, 0, 0, 0}; + mjuu_frameaccumChild(frame->pos, frame->quat, pos, qunit); + + // rotate dir + mjuu_rotVecQuat(dir, dir, frame->quat); } // normalize direction, make sure it is not zero diff --git a/test/user/user_objects_test.cc b/test/user/user_objects_test.cc index a23d937e..6cd6fabf 100644 --- a/test/user/user_objects_test.cc +++ b/test/user/user_objects_test.cc @@ -2086,7 +2086,6 @@ TEST_F(MujocoTest, Frame) { - )"; constexpr mjtNum eps = 1e-14; std::array error; @@ -2142,6 +2141,34 @@ TEST_F(MujocoTest, Frame) { mj_deleteData(d); } +TEST_F(MujocoTest, FrameTransformsLight) { + static constexpr char xml[] = R"( + + + + + + + + )"; + std::array error; + mjModel* m = LoadModelFromString(xml, error.data(), error.size()); + EXPECT_THAT(m, testing::NotNull()) << error.data(); + EXPECT_EQ(m->nlight, 1); + + constexpr mjtNum eps = 1e-14; + EXPECT_NEAR(m->light_pos[0], -mju_sqrt(.5), eps); + EXPECT_NEAR(m->light_pos[1], 0, eps); + EXPECT_NEAR(m->light_pos[2], 1 + mju_sqrt(.5), eps); + + EXPECT_NEAR(m->light_dir[0], 0, eps); + EXPECT_NEAR(m->light_dir[1], 0, eps); + EXPECT_NEAR(m->light_dir[2], -1, eps); + + mj_deleteModel(m); +} + + // ------------- test bvh ------------------------------------------------------ TEST_F(MujocoTest, RobustBVH) { static constexpr char xml1[] = R"( From cf3e60f9a42aa20c5546993045e8d521f728fcf6 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 21 Aug 2024 02:53:03 -0700 Subject: [PATCH 12/26] Fix load time reporting in simulate. Before this change, a model which took more than 1/4 seconds to load caused the simulation to pause, which is a behavior that is only supposed to occur for models that loaded with a warning. PiperOrigin-RevId: 665786861 Change-Id: Ie3a6c8d18a6abbb425a0bbcc5a8e55d5fd7a0c72 --- simulate/main.cc | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/simulate/main.cc b/simulate/main.cc index 15316383..ac4d7b36 100644 --- a/simulate/main.cc +++ b/simulate/main.cc @@ -245,13 +245,6 @@ mjModel* LoadModel(const char* file, mj::Simulate& sim) { auto load_interval = mj::Simulate::Clock::now() - load_start; double load_seconds = Seconds(load_interval).count(); - // if no error and load took more than 1/4 seconds, report load time - if (!loadError[0] && load_seconds > 0.25) { - mju::sprintf_arr(loadError, "Model loaded in %.2g seconds", load_seconds); - } - - mju::strcpy_arr(sim.load_error, loadError); - if (!mnew) { std::printf("%s\n", loadError); return nullptr; @@ -264,6 +257,13 @@ mjModel* LoadModel(const char* file, mj::Simulate& sim) { sim.run = 0; } + // if no error and load took more than 1/4 seconds, report load time + if (!loadError[0] && load_seconds > 0.25) { + mju::sprintf_arr(loadError, "Model loaded in %.2g seconds", load_seconds); + } + + mju::strcpy_arr(sim.load_error, loadError); + return mnew; } From f986a52e1e6873fbe82ef660e70328b79b856686 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Wed, 21 Aug 2024 03:05:17 -0700 Subject: [PATCH 13/26] Compute witness points from EPA, and use witness points to determine pos and dir for contact. PiperOrigin-RevId: 665791337 Change-Id: If33ca1f455804c4cda534508803f5c3e99c3c52a --- src/engine/engine_collision_gjk.c | 110 ++++++++++------------- test/engine/engine_collision_gjk_test.cc | 1 + 2 files changed, 49 insertions(+), 62 deletions(-) diff --git a/src/engine/engine_collision_gjk.c b/src/engine/engine_collision_gjk.c index 2118d67a..daab5889 100644 --- a/src/engine/engine_collision_gjk.c +++ b/src/engine/engine_collision_gjk.c @@ -96,8 +96,9 @@ static int newVertex(Polytope* pt, const mjtNum v1[3], const mjtNum v2[3]); static void attachFace(Polytope* pt, int v1, int v2, int v3); // returns the penetration depth (negative distance) of the convex objects +// witness points are stored in x1 and x2 static mjtNum epa(const mjCCDConfig* config, Polytope* pt, - mjCCDObj* obj1, mjCCDObj* obj2, Face* nearest); + mjCCDObj* obj1, mjCCDObj* obj2, mjtNum x1[3], mjtNum x2[3]); // internal data structure for the returning simplex from GJK typedef struct { @@ -858,10 +859,37 @@ static void addEdgeIfUnique(Horizon* h, int v1, int v2) { #undef mjMINCAP +// recover witness points from EPA polytope +static void epa_witness(const Polytope* pt, int index, mjtNum x1[3], mjtNum x2[3]) { + Face* face = &pt->faces[index]; + int s1 = face->verts[0], s2 = face->verts[1], s3 = face->verts[2]; + + // run S2D to get barycentric coordinates of witness point + // witness point is guaranteed to be an internal point of face + mjtNum simplex[9], lambda[4]; + mju_copy3(simplex, pt->verts[s1].v); + mju_copy3(simplex + 3, pt->verts[s2].v); + mju_copy3(simplex + 6, pt->verts[s3].v); + S2D(lambda, simplex); + + // face on geom 1 + mjtNum simplex1[9]; + mju_copy3(simplex1, pt->verts[s1].v1); + mju_copy3(simplex1 + 3, pt->verts[s2].v1); + mju_copy3(simplex1 + 6, pt->verts[s3].v1); + lincomb(x1, lambda, simplex1, 3); + + // face on geom 2 + mjtNum simplex2[9]; + mju_copy3(simplex2, pt->verts[s1].v2); + mju_copy3(simplex2 + 3, pt->verts[s2].v2); + mju_copy3(simplex2 + 6, pt->verts[s3].v2); + lincomb(x2, lambda, simplex2, 3); +} // returns the penetration depth (negative distance) of the convex objects static mjtNum epa(const mjCCDConfig* config, Polytope* pt, - mjCCDObj* obj1, mjCCDObj* obj2, Face* nearest) { + mjCCDObj* obj1, mjCCDObj* obj2, mjtNum x1[3], mjtNum x2[3]) { mjtNum dist = mjMAXVAL; int index; Horizon h; @@ -913,8 +941,7 @@ static mjtNum epa(const mjCCDConfig* config, Polytope* pt, h.n = 0; // clear horizon } mju_free(h.edges); - nearest->dist = dist; - mju_copy3(nearest->n, pt->faces[index].n); + epa_witness(pt, index, x1, x2); return dist; } @@ -922,7 +949,7 @@ static mjtNum epa(const mjCCDConfig* config, Polytope* pt, // runs both GJK and EPA (if needed) static mjtNum _gjk_epa(const mjCCDConfig* config, mjCCDObj* obj1, mjCCDObj* obj2, Polytope* pt, - Face* nearest) { + mjtNum x1[3], mjtNum x2[3]) { Simplex simplex1, simplex2; mjtNum dist = _gjk(config, obj1, obj2, &simplex1, &simplex2); @@ -938,8 +965,8 @@ static mjtNum _gjk_epa(const mjCCDConfig* config, mjCCDObj* obj1, mjCCDObj* obj2 // simplex not on boundary (objects are penetrating) if (ret) { - epa(config, pt, obj1, obj2, nearest); - return -nearest->dist; + dist = epa(config, pt, obj1, obj2, x1, x2); + return -dist; } return 0; } @@ -948,76 +975,35 @@ static mjtNum _gjk_epa(const mjCCDConfig* config, mjCCDObj* obj1, mjCCDObj* obj2 // --------------------------- LibCCD Compatibility Layer ----------------------------------------- -static int posCompare(const void *a, const void *b) { - Vertex *v1, *v2; - v1 = *(Vertex**) a; - v2 = *(Vertex**) b; - - if (v1->dist == v2->dist) { - return 0; - } else if (v1->dist < v2->dist) { - return -1; - } else { - return 1; - } -} - - - -// computes the position of contact in the same manner as LibCCD -static int computePos(const Polytope* pt, mjtNum pos[3]) { - Vertex** vs; - int len = pt->nverts; - mjtNum scale = 0; - - vs = (Vertex**) mju_malloc(len * sizeof(Vertex*)); - if (vs == NULL) return -1; - - for (int i = 0; i < len; i++) { - vs[i] = pt->verts + i; - } - - qsort(vs, len, sizeof(Vertex*), posCompare); - - mju_zero3(pos); - if (len % 2 == 1) len++; - - // average out the vertices of the polytope - for (int i = 0; i < len / 2; i++) { - mju_add3(pos, pos, vs[i]->v1); - mju_add3(pos, pos, vs[i]->v2); - scale += 2; - } - mju_scl3(pos, pos, 1 / scale); - - mju_free(vs); - return 0; -} - - - // Penetration function with same signature as LibCCD's ccdMPRPenetration and ccdGJKPenetration int mj_gjkPenetration(const void *obj1, const void *obj2, const ccd_t *ccd, ccd_real_t *depth, ccd_vec3_t *dir, ccd_vec3_t *pos) { Polytope pt; initPolytope(&pt); - Face nearest; mjCCDConfig config; mjCCDObj* o1 = (mjCCDObj*) obj1; mjCCDObj* o2 = (mjCCDObj*) obj2; - nearest.n[1] = 34; - o1->center(o1->x0, o1); o2->center(o2->x0, o2); config.max_iterations = ccd->max_iterations; config.tolerance = ccd->mpr_tolerance; - mjtNum dist = _gjk_epa(&config, o1, o2, &pt, &nearest); + mjtNum x1[3], x2[3]; + mjtNum dist = _gjk_epa(&config, o1, o2, &pt, x1, x2); if (dist < 0) { - if (depth) *depth = nearest.dist; - if (dir) mju_copy3(dir->v, nearest.n); - if (pos) computePos(&pt, pos->v); + if (depth) *depth = -dist; + if (dir) { + mjtNum d[3]; + mju_sub3(d, x1, x2); + mju_normalize3(d); + mju_copy3(dir->v, d); + } + if (pos) { + mju_scl3(x1, x1, 0.5); + mju_scl3(x2, x2, 0.5); + mju_add3(pos->v, x1, x2); + } } else { if (depth) *depth = 0; if (dir) mju_zero3(dir->v); diff --git a/test/engine/engine_collision_gjk_test.cc b/test/engine/engine_collision_gjk_test.cc index f7a30e1b..da1d1602 100644 --- a/test/engine/engine_collision_gjk_test.cc +++ b/test/engine/engine_collision_gjk_test.cc @@ -70,6 +70,7 @@ mjtNum run_gjkPenetration(mjModel* m, mjData* d, int g1, int g2, mjCCDObj obj2 = {m, d, g2, -1, -1, -1, -1, 0, {1, 0, 0, 0}, {0, 0, 0}, mjc_center, mjc_support}; ccd_t ccd; + // CCD_INIT(&ccd); // uncomment to run ccdMPRPenetration ccd.mpr_tolerance = kTolerance; ccd.epa_tolerance = kTolerance; ccd.max_iterations = kMaxIterations; From 806b8c8e2e61fc9edfa69b4ed19ac2a2a4ca3d58 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 21 Aug 2024 03:45:40 -0700 Subject: [PATCH 14/26] Add support for actuatorfrc and jointactfrc sensors in MJX. PiperOrigin-RevId: 665804850 Change-Id: I99371fc4da055382f1564b2c67d25f7b30a0a0cf --- doc/changelog.rst | 13 +++++++------ mjx/mujoco/mjx/_src/sensor.py | 28 +++++++++++++++++++++++++++- mjx/mujoco/mjx/_src/types.py | 4 ++++ mjx/mujoco/mjx/test_data/sensor.xml | 6 ++++++ 4 files changed, 44 insertions(+), 7 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 9b47e771..695122ee 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -21,21 +21,22 @@ MJX ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, ``CLOCK``. 7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. -8. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. -9. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. -10. Added support for :ref:`implicitfast integration` for all cases except +8. Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``. +9. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. +10. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. +11. Added support for :ref:`implicitfast integration` for all cases except :doc:`fluid drag `. Bug fixes ^^^^^^^^^ -11. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, +12. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, contribution by :github:user:`michael-ahn`). -12. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit +13. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit integrators, wrong derivatives would be computed. Python bindings ^^^^^^^^^^^^^^^ -13. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). +14. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). Version 3.2.2 (Aug 8, 2024) diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index 77f6a3aa..4e704472 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -263,4 +263,30 @@ def sensor_acc(m: Model, d: Data) -> Data: if m.opt.disableflags & DisableBit.SENSOR: return d - return d + stage_acc = m.sensor_needstage == mujoco.mjtStage.mjSTAGE_ACC + sensors, adrs = [], [] + + for sensor_type in set(m.sensor_type[stage_acc]): + idx = m.sensor_type == sensor_type + objid = m.sensor_objid[idx] + adr = m.sensor_adr[idx] + + if sensor_type == SensorType.ACTUATORFRC: + sensor = d.actuator_force[objid] + elif sensor_type == SensorType.JOINTACTFRC: + sensor = d.qfrc_actuator[m.jnt_dofadr[objid]] + else: + # TODO(taylorhowell): raise error after adding sensor check to io.py + continue # unsupported sensor type + + sensors.append(sensor) + adrs.append(adr) + + if not adrs: + return d + + sensordata = d.sensordata.at[np.concatenate(adrs)].set( + jp.concatenate(sensors) + ) + + return d.replace(sensordata=sensordata) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 4a14db17..fd3aa362 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -309,6 +309,8 @@ class SensorType(enum.IntEnum): JOINTVEL: joint velocity ACTUATORVEL: actuator velocity BALLANGVEL: ball joint angular velocity + ACTUATORFRC: scalar actuator force + JOINTACTFRC: scalar actuator force, measured at the joint """ MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION @@ -325,6 +327,8 @@ class SensorType(enum.IntEnum): JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL ACTUATORVEL = mujoco.mjtSensor.mjSENS_ACTUATORVEL BALLANGVEL = mujoco.mjtSensor.mjSENS_BALLANGVEL + ACTUATORFRC = mujoco.mjtSensor.mjSENS_ACTUATORFRC + JOINTACTFRC = mujoco.mjtSensor.mjSENS_JOINTACTFRC class ObjType(PyTreeNode): diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml index 04924952..80285609 100644 --- a/mjx/mujoco/mjx/test_data/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor.xml @@ -18,6 +18,8 @@ -actuatorvel -ballangvel * acceleration/force-dependent sensors: +-actuatorfrc +-jointactfrc --> @@ -81,14 +83,17 @@ + + + @@ -97,6 +102,7 @@ + From b42780a5f4726333d815f28be086589637b75a06 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Wed, 21 Aug 2024 16:20:39 -0700 Subject: [PATCH 15/26] Add subtree_vel function to MJX. This function matches mj_subtreeVel. PiperOrigin-RevId: 666078305 Change-Id: Ia0ecbf8279e914f49ee60b862541d697e8d4abb1 --- mjx/mujoco/mjx/__init__.py | 1 + mjx/mujoco/mjx/_src/smooth.py | 93 ++++++++++++++++++++++++++++++ mjx/mujoco/mjx/_src/smooth_test.py | 19 ++++++ 3 files changed, 113 insertions(+) diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index a4a8a5da..d272dd13 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -43,6 +43,7 @@ from mujoco.mjx._src.smooth import crb from mujoco.mjx._src.smooth import factor_m from mujoco.mjx._src.smooth import kinematics from mujoco.mjx._src.smooth import rne +from mujoco.mjx._src.smooth import subtree_vel from mujoco.mjx._src.smooth import tendon from mujoco.mjx._src.smooth import transmission from mujoco.mjx._src.solver import solve diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index ed6f3dce..96918c6a 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -427,6 +427,99 @@ def com_vel(m: Model, d: Data) -> Data: return d +def subtree_vel(m: Model, d: Data) -> Data: + """Subtree linear velocity and angular momentum.""" + + # bodywise quantities + def _forward(cvel, xipos, ximat, subtree_com_root, mass, inertia): + ang, lin = jp.split(cvel, 2) + + # update linear velocity + lin = lin - jp.cross(xipos - subtree_com_root, ang) + + subtree_linvel = mass * lin + subtree_angmom = inertia * ximat @ ximat.T @ ang + body_vel = jp.concatenate([ang, lin]) + + return body_vel, subtree_linvel, subtree_angmom + + body_vel, subtree_linvel, subtree_angmom = jax.vmap(_forward)( + d.cvel, + d.xipos, + d.ximat, + d.subtree_com[m.body_rootid], + m.body_mass, + m.body_inertia, + ) + + # sum body linear momentum recursively up the kinematic tree + subtree_linvel = scan.body_tree( + m, + lambda x, y: y if x is None else x + y, + 'bb', + 'b', + subtree_linvel, + reverse=True, + ) + + subtree_linvel /= jp.maximum(mujoco.mjMINVAL, m.body_subtreemass)[:, None] + + def _subtree_angmom( + carry, + angmom, + com, + com_parent, + linvel, + linvel_parent, + subtreemass, + xipos, + vel, + mass, + mask, + ): + + def _momentum(x0, x1, v0, v1, m): + dx = x0 - x1 + dv = v0 - v1 + dp = dv * m + return jp.cross(dx, dp) + + # momentum wrt current body + mom = mask * _momentum(xipos, com, vel[3:], linvel, mass) + + # momentum wrt parent + mom_parent = mask * _momentum( + com, com_parent, linvel, linvel_parent, subtreemass + ) + + if carry is None: + return angmom + mom, mom_parent + else: + angmom_child, mom_parent_child = carry + return angmom + mom + angmom_child + mom_parent_child, mom_parent + + + subtree_angmom, _ = scan.body_tree( + m, + _subtree_angmom, + 'bbbbbbbbbb', + 'bb', + subtree_angmom, + d.subtree_com, + d.subtree_com[m.body_parentid], + subtree_linvel, + subtree_linvel[m.body_parentid], + m.body_subtreemass, + d.xipos, + body_vel, + m.body_mass, + jp.ones(m.nbody).at[0].set(0), + reverse=True, + ) + + return d.replace(subtree_linvel=subtree_linvel, subtree_angmom=subtree_angmom) + + def rne(m: Model, d: Data) -> Data: """Computes inverse dynamics using the recursive Newton-Euler algorithm.""" # forward scan over tree: accumulate link center of mass acceleration diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index f0efc07f..c8e5f162 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -172,6 +172,25 @@ class SmoothTest(absltest.TestCase): _assert_attr_eq(d, dx, 'actuator_length') _assert_attr_eq(d, dx, 'actuator_moment') + def test_subtree_vel(self): + """Tests MJX subtree_vel function matches MuJoCo mj_subtreeVel.""" + + m = test_util.load_test_file('humanoid/humanoid.xml') + d = mujoco.MjData(m) + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) + mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) + + # subtree velocity + mujoco.mj_subtreeVel(m, d) + dx = jax.jit(mjx.subtree_vel)(mx, dx) + + _assert_attr_eq(d, dx, 'subtree_linvel') + _assert_attr_eq(d, dx, 'subtree_angmom') + if __name__ == '__main__': absltest.main() From b2174a7ec30dbd0c8dc554699e9d3c4dc4731233 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Thu, 22 Aug 2024 02:20:20 -0700 Subject: [PATCH 16/26] Add framequat sensor to MJX. PiperOrigin-RevId: 666254481 Change-Id: I50be8574e1e4677ef5dc23a2bdcbc296f301d7fc --- doc/changelog.rst | 4 +- mjx/mujoco/mjx/_src/sensor.py | 94 ++++++++++++++++++++--------- mjx/mujoco/mjx/_src/types.py | 2 + mjx/mujoco/mjx/test_data/sensor.xml | 4 ++ 4 files changed, 72 insertions(+), 32 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 695122ee..195d20cd 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -18,8 +18,8 @@ MJX ^^^ 5. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`). 6. Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, - ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``SUBTREECOM``, - ``CLOCK``. + ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, + ``SUBTREECOM``, ``CLOCK``. 7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. 8. Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``. 9. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py index 4e704472..d2ad1497 100644 --- a/mjx/mujoco/mjx/_src/sensor.py +++ b/mjx/mujoco/mjx/_src/sensor.py @@ -61,6 +61,9 @@ def sensor_pos(m: Model, d: Data) -> Data: for sensor_type in set(m.sensor_type[stage_pos]): idx = m.sensor_type == sensor_type objid = m.sensor_objid[idx] + objtype = m.sensor_objtype[idx] + refid = m.sensor_refid[idx] + reftype = m.sensor_reftype[idx] adr = m.sensor_adr[idx] if sensor_type == SensorType.MAGNETOMETER: @@ -113,7 +116,6 @@ def sensor_pos(m: Model, d: Data) -> Data: return sensor[:2] - refid = m.sensor_refid[idx] sensorsize = m.cam_sensorsize[refid] intrinsic = m.cam_intrinsic[refid] fovy = m.cam_fovy[refid] @@ -131,15 +133,15 @@ def sensor_pos(m: Model, d: Data) -> Data: elif sensor_type == SensorType.RANGEFINDER: site_bodyid = m.site_bodyid[objid] for sid in set(site_bodyid): - id_ = sid == site_bodyid - objid_ = objid[id_] - site_xpos = d.site_xpos[objid_] - site_mat = d.site_xmat[objid_].reshape((-1, 9))[:, np.array([2, 5, 8])] + idxs = sid == site_bodyid + objids = objid[idxs] + site_xpos = d.site_xpos[objids] + site_mat = d.site_xmat[objids].reshape((-1, 9))[:, np.array([2, 5, 8])] sensor, _ = jax.vmap( ray.ray, in_axes=(None, None, 0, 0, None, None, None) )(m, d, site_xpos, site_mat, (), True, sid) sensors.append(sensor) - adrs.append(adr[id_]) + adrs.append(adr[idxs]) continue # avoid adding to sensors/adrs list a second time elif sensor_type == SensorType.JOINTPOS: sensor = d.qpos[m.jnt_qposadr[objid]] @@ -155,23 +157,19 @@ def sensor_pos(m: Model, d: Data) -> Data: def _framepos(xpos, xpos_ref, xmat_ref, refid): return jp.where(refid == -1, xpos, xmat_ref.T @ (xpos - xpos_ref)) - objtype = m.sensor_objtype[idx] - reftype = m.sensor_reftype[idx] - refid = m.sensor_refid[idx] - # evaluate for valid object and reference object type pairs for ot, rt in set(zip(objtype, reftype)): - id_ = (objtype == ot) & (reftype == rt) - refid_ = refid[id_] + idxt = (objtype == ot) & (reftype == rt) + refidt = refid[idxt] xpos, _ = objtype_data[ot] xpos_ref, xmat_ref = objtype_data[rt] - xpos = xpos[objid[id_]] - xpos_ref = xpos_ref[refid_] - xmat_ref = xmat_ref[refid_] - sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refid_) - adr_ = adr[id_, None] + np.arange(3)[None] + xpos = xpos[objid[idxt]] + xpos_ref = xpos_ref[refidt] + xmat_ref = xmat_ref[refidt] + sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refidt) + adrt = adr[idxt, None] + np.arange(3)[None] sensors.append(sensor.reshape(-1)) - adrs.append(adr_.reshape(-1)) + adrs.append(adrt.reshape(-1)) continue # avoid adding to sensors/adrs list a second time elif sensor_type in frame_axis: @@ -179,22 +177,58 @@ def sensor_pos(m: Model, d: Data) -> Data: axis = xmat[:, frame_axis[sensor_type]] return jp.where(refid == -1, axis, xmat_ref.T @ axis) - objtype = m.sensor_objtype[idx] - reftype = m.sensor_reftype[idx] - refid = m.sensor_refid[idx] + # evaluate for valid object and reference object type pairs + for ot, rt in set(zip(objtype, reftype)): + idxt = (objtype == ot) & (reftype == rt) + refidt = refid[idxt] + _, xmat = objtype_data[ot] + _, xmat_ref = objtype_data[rt] + xmat = xmat[objid[idxt]] + xmat_ref = xmat_ref[refidt] + sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refidt) + adrt = adr[idxt, None] + np.arange(3)[None] + sensors.append(sensor.reshape(-1)) + adrs.append(adrt.reshape(-1)) + continue # avoid adding to sensors/adrs list a second time + elif sensor_type == SensorType.FRAMEQUAT: + + def _quat(otype, oid): + if otype == ObjType.XBODY: + return d.xquat[oid] + elif otype == ObjType.BODY: + return jax.vmap(math.quat_mul)(d.xquat[oid], m.body_iquat[oid]) + elif otype == ObjType.GEOM: + return jax.vmap(math.quat_mul)( + d.xquat[m.geom_bodyid[oid]], m.geom_quat[oid] + ) + elif otype == ObjType.SITE: + return jax.vmap(math.quat_mul)( + d.xquat[m.site_bodyid[oid]], m.site_quat[oid] + ) + elif otype == ObjType.CAMERA: + return jax.vmap(math.quat_mul)( + d.xquat[m.cam_bodyid[oid]], m.cam_quat[oid] + ) + elif otype == ObjType.UNKNOWN: + return jp.tile(jp.array([1.0, 0.0, 0.0, 0.0]), (oid.size, 1)) + else: + raise ValueError(f'Unknown object type: {otype}') # evaluate for valid object and reference object type pairs for ot, rt in set(zip(objtype, reftype)): - id_ = (objtype == ot) & (reftype == rt) - refid_ = refid[id_] - _, xmat = objtype_data[ot] - _, xmat_ref = objtype_data[rt] - xmat = xmat[objid[id_]] - xmat_ref = xmat_ref[refid_] - sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refid_) - adr_ = adr[id_, None] + np.arange(3)[None] + idxt = (objtype == ot) & (reftype == rt) + objidt = objid[idxt] + refidt = refid[idxt] + quat = _quat(ot, objidt) + refquat = _quat(rt, refidt) + sensor = jax.vmap( + lambda q, r, rid: jp.where( + rid == -1, q, math.quat_mul(math.quat_inv(r), q) + ) + )(quat, refquat, refidt) + adrt = adr[idxt, None] + np.arange(4)[None] sensors.append(sensor.reshape(-1)) - adrs.append(adr_.reshape(-1)) + adrs.append(adrt.reshape(-1)) continue # avoid adding to sensors/adrs list a second time elif sensor_type == SensorType.SUBTREECOM: sensor = d.subtree_com[objid].reshape(-1) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index fd3aa362..eaf54fd4 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -304,6 +304,7 @@ class SensorType(enum.IntEnum): FRAMEXAXIS: frame x-axis FRAMEYAXIS: frame y-axis FRAMEZAXIS: frame z-axis + FRAMEQUAT: frame orientation, represented as quaternion SUBTREECOM: subtree centor of mass CLOCK: simulation time JOINTVEL: joint velocity @@ -322,6 +323,7 @@ class SensorType(enum.IntEnum): FRAMEXAXIS = mujoco.mjtSensor.mjSENS_FRAMEXAXIS FRAMEYAXIS = mujoco.mjtSensor.mjSENS_FRAMEYAXIS FRAMEZAXIS = mujoco.mjtSensor.mjSENS_FRAMEZAXIS + FRAMEQUAT = mujoco.mjtSensor.mjSENS_FRAMEQUAT SUBTREECOM = mujoco.mjtSensor.mjSENS_SUBTREECOM CLOCK = mujoco.mjtSensor.mjSENS_CLOCK JOINTVEL = mujoco.mjtSensor.mjSENS_JOINTVEL diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml index 80285609..5a636e31 100644 --- a/mjx/mujoco/mjx/test_data/sensor.xml +++ b/mjx/mujoco/mjx/test_data/sensor.xml @@ -11,6 +11,7 @@ -framexaxis -frameyaxis -framezaxis +-framequat -subtreecom -clock * velocity-dependent sensors: @@ -50,6 +51,7 @@ + @@ -91,6 +93,7 @@ + @@ -109,6 +112,7 @@ + From f75ead2203c9204d45b50ec0fb32960c25cdb7fd Mon Sep 17 00:00:00 2001 From: Nimrod Gileadi Date: Thu, 22 Aug 2024 03:21:19 -0700 Subject: [PATCH 17/26] Improve the build instructions for mac. The app bundle doesn't include mujoco.framework, so the instructions were a bit stale. I verified that with these new instructions I'm able to build the Python bindings from source with a prebuilt MuJoCo binary. Fixes #313. PiperOrigin-RevId: 666272041 Change-Id: I8bcf14c8ab6268cdaa617a8b5940b155ef64a8c5 --- doc/programming/index.rst | 20 ++++++-------------- doc/python.rst | 21 +++++++++++---------- 2 files changed, 17 insertions(+), 24 deletions(-) diff --git a/doc/programming/index.rst b/doc/programming/index.rst index a7424422..03916e22 100644 --- a/doc/programming/index.rst +++ b/doc/programming/index.rst @@ -76,24 +76,16 @@ directory; it contains error and warning messages, and can be deleted at any tim After verifying that the simulator works, you may also want to re-compile the code samples to ensure that you have a working development environment. We provide a cross-platform `CMake -`_ setup that can be used to build sample +`__ setup that can be used to build sample applications independently of the MuJoCo library itself. On macOS, the DMG disk image contains ``MuJoCo.app``, which you can double-click to launch the ``simulate`` GUI. You can -also drag ``MuJoCo.app`` into the ``/Application`` on your system, as you would to install any other app. While -``MuJoCo.app`` may look like a file, it is in fact an `Application Bundle `_, which is a directory that contains executable binaries for all of MuJoCo's sample applications, along with -an embedded `framework -`_, -which is a subdirectory containing the MuJoCo dynamic library and all of its public headers. In other words, -``MuJoCo.app`` contains all the same files that are shipped in the archive on Windows and Linux. To see this, right -click (or control-click) on ``MuJoCo.app`` and click "Show Package Contents". - -As mentioned above, ``mujoco.framework`` contains the library and headers that are necessary to build any application -that depends on MuJoCo. If you are using Xcode, you can import it as a framework dependency on your project. (This also +also drag ``MuJoCo.app`` into the ``/Application`` on your system, as you would to install any other app. As well as the +``MuJoCo.app`` `Application Bundle `__, the DMG includes the ``mujoco.framework`` subdirectory containing the MuJoCo dynamic library and all of +its public headers. If you are using Xcode, you can import it as a framework dependency on your project. (This also works for Swift projects without any modification). If you are building manually, you can use ``-F`` and -``-framework mujoco`` to specify the header search path and the library search path respectively. The macOS Makefile -provides an example for this. +``-framework mujoco`` to specify the header search path and the library search path respectively. .. _inBuild: diff --git a/doc/python.rst b/doc/python.rst index be1127fd..2454c075 100644 --- a/doc/python.rst +++ b/doc/python.rst @@ -522,8 +522,8 @@ Building from source 1. Make sure you have CMake and a C++17 compiler installed. 2. Download the `latest binary release `__ - from GitHub. On macOS, the download corresponds to a DMG file from which you - can drag ``MuJoCo.app`` into your ``/Applications`` folder. + from GitHub. On macOS, the download corresponds to a DMG file which you can mount by + double-clicking or running ``hdiutil attach ``. 3. Clone the entire ``mujoco`` repository from GitHub and ``cd`` into the python directory: @@ -545,7 +545,6 @@ Building from source .. code-block:: shell - cd python bash make_sdist.sh The ``make_sdist.sh`` script generates additional C++ header files that are @@ -556,19 +555,21 @@ Building from source 6. Use the generated source distribution to build and install the bindings. You'll need to specify the path to the MuJoCo library you downloaded earlier - in the ``MUJOCO_PATH`` environment variable. + in the ``MUJOCO_PATH`` environment variable, and the path to the MuJoCo + plugin directory in the ``MUJOCO_PLUGIN_PATH`` environment variable. .. note:: - For macOS, this can be the path to a directory that contains the - ``mujoco.framework``. In particular, you can set - ``MUJOCO_PATH=/Applications/MuJoCo.app`` if you installed MuJoCo as - suggested in step 1. + For macOS, the files need to be extracted from the DMG. + Once you mounted it as in step 2, the ``mujoco.framework`` directory can be found in ``/Volumes/MuJoCo``, + and the plugins directory can be found in ``/Volumes/MuJoCo/MuJoCo.app/Contents/MacOS/mujoco_plugin``. + Those two directories can be copied out somewhere convenient, or you can use + ``MUJOCO_PATH=/Volumes/MuJoCo MUJOCO_PLUGIN_PATH=/Volumes/MuJoCo/MuJoCo.app/Contents/MacOS/mujoco_plugin``. .. code-block:: shell cd dist - MUJOCO_PATH=/PATH/TO/MUJOCO - MUJOCO_PLUGIN_PATH=/PATH/TO/MUJOCO_PLUGIN + MUJOCO_PATH=/PATH/TO/MUJOCO \ + MUJOCO_PLUGIN_PATH=/PATH/TO/MUJOCO_PLUGIN \ pip install mujoco-x.y.z.tar.gz The Python bindings should now be installed! To check that they've been From a2649d6cec582009f429d5126e69a68cf6a8ede4 Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Thu, 22 Aug 2024 03:26:04 -0700 Subject: [PATCH 18/26] Add efc_margin to MJX. PiperOrigin-RevId: 666273240 Change-Id: I12d6cc0d3edd3c3faff5a876335910b42e9cc9bf --- mjx/mujoco/mjx/_src/constraint.py | 61 ++++++++++++++++++++------ mjx/mujoco/mjx/_src/constraint_test.py | 28 +++++++++++- mjx/mujoco/mjx/_src/io.py | 2 + mjx/mujoco/mjx/_src/types.py | 2 + 4 files changed, 78 insertions(+), 15 deletions(-) diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index 8f0bbf92..40111298 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -44,6 +44,7 @@ class _Efc(PyTreeNode): invweight: jax.Array solref: jax.Array solimp: jax.Array + margin: jax.Array def _kbi( @@ -119,7 +120,7 @@ def _efc_equality_connect(m: Model, d: Data) -> Optional[_Efc]: pos_imp = math.norm(pos) invweight = m.body_invweight0[obj1id, 0] + m.body_invweight0[obj2id, 0] - return _row(j, pos, pos_imp, invweight, solref, solimp) + return _row(j, pos, pos_imp, invweight, solref, solimp, jp.zeros_like(pos)) args = (m.eq_obj1id, m.eq_obj2id, m.eq_data, m.eq_solref, m.eq_solimp) args = jax.tree_util.tree_map(lambda x: x[eq_id], args) @@ -165,7 +166,7 @@ def _efc_equality_weld(m: Model, d: Data) -> Optional[_Efc]: invweight = m.body_invweight0[obj1id] + m.body_invweight0[obj2id] invweight = jp.repeat(invweight, 3, axis=0) - return _row(j, pos, pos_imp, invweight, solref, solimp) + return _row(j, pos, pos_imp, invweight, solref, solimp, jp.zeros_like(pos)) args = (m.eq_obj1id, m.eq_obj2id, m.eq_data, m.eq_solref, m.eq_solimp) args = jax.tree_util.tree_map(lambda x: x[eq_id], args) @@ -194,7 +195,7 @@ def _efc_equality_joint(m: Model, d: Data) -> Optional[_Efc]: invweight = m.dof_invweight0[dofadr1] invweight += m.dof_invweight0[dofadr2] * (obj2id > -1) - return _row(j, pos, pos, invweight, solref, solimp) + return _row(j, pos, pos, invweight, solref, solimp, jp.zeros_like(pos)) args = (m.eq_obj1id, m.eq_obj2id, m.eq_data, m.eq_solref, m.eq_solimp) args = jax.tree_util.tree_map(lambda x: x[eq_id], args) @@ -232,7 +233,7 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]: deriv = jp.dot(data[1:5], dif_power[:4] * jp.arange(1, 5)) * (obj2id > -1) j = jac1 + jac2 * -deriv - return _row(j, pos, pos, invweight, solref, solimp) + return _row(j, pos, pos, invweight, solref, solimp, jp.zeros_like(pos)) inv1, inv2 = m.tendon_invweight0[obj1id], m.tendon_invweight0[obj2id] jac1, jac2 = d.ten_J[obj1id], d.ten_J[obj2id] @@ -267,7 +268,9 @@ def _efc_limit_ball(m: Model, d: Data) -> Optional[_Efc]: j = jp.zeros(m.nv).at[jp.arange(3) + dofadr].set(-axis) invweight = m.dof_invweight0[dofadr] - return _row(j * active, pos * active, pos, invweight, solref, solimp) + return _row( + j * active, pos * active, pos, invweight, solref, solimp, jnt_margin + ) args = (m.jnt_qposadr, m.jnt_dofadr, m.jnt_range, m.jnt_margin, m.jnt_solref) args += (m.jnt_solimp,) @@ -294,7 +297,9 @@ def _efc_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]: j = jp.zeros(m.nv).at[dofadr].set((dist_min < dist_max) * 2 - 1) invweight = m.dof_invweight0[dofadr] - return _row(j * active, pos * active, pos, invweight, solref, solimp) + return _row( + j * active, pos * active, pos, invweight, solref, solimp, jnt_margin + ) args = (m.jnt_qposadr, m.jnt_dofadr, m.jnt_range, m.jnt_margin, m.jnt_solref) args += (m.jnt_solimp,) @@ -328,7 +333,7 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]: active = pos < 0 j = jax.vmap(jp.multiply)(j, ((dist_min < dist_max) * 2 - 1) * active) - return jax.vmap(_row)(j, pos * active, pos, invweight, solref, solimp) + return jax.vmap(_row)(j, pos * active, pos, invweight, solref, solimp, margin) def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]: @@ -349,7 +354,15 @@ def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]: j = (c.frame @ (jac2p - jac1p).T)[0] invweight = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0] - return _row(j * active, pos * active, pos, invweight, c.solref, c.solimp) + return _row( + j * active, + pos * active, + pos, + invweight, + c.solref, + c.solimp, + c.includemargin, + ) contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact) @@ -385,7 +398,15 @@ def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]: invweight = invweight + fri[0] * fri[0] * invweight invweight = invweight * 2 * fri[0] * fri[0] / m.opt.impratio - return _row(j * active, pos * active, pos, invweight, c.solref, c.solimp) + return _row( + j * active, + pos * active, + pos, + invweight, + c.solref, + c.solimp, + c.includemargin, + ) contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact) # concatenate to drop row grouping @@ -421,7 +442,15 @@ def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]: invweight = jp.concatenate((invweight, invweight[1] * fri)) pos_aref = jp.zeros(condim).at[0].set(pos) - return _row(j * active, pos_aref * active, pos, invweight, solref, c.solimp) + return _row( + j * active, + pos_aref * active, + pos, + invweight, + solref, + c.solimp, + c.includemargin, + ) contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact) # concatenate to drop row grouping @@ -529,7 +558,9 @@ def make_constraint(m: Model, d: Data) -> Data: if not efcs: z = jp.empty(0) d = d.replace(efc_J=jp.empty((0, m.nv))) - d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z, efc_pos=z) + d = d.replace( + efc_D=z, efc_aref=z, efc_frictionloss=z, efc_pos=z, efc_margin=z + ) return d efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs) @@ -539,10 +570,12 @@ def make_constraint(m: Model, d: Data) -> Data: k, b, imp = _kbi(m, efc.solref, efc.solimp, efc.pos_imp) r = jp.maximum(efc.invweight * (1 - imp) / imp, mujoco.mjMINVAL) aref = -b * (efc.J @ d.qvel) - k * imp * efc.pos_aref - return aref, r, efc.pos_aref + return aref, r, efc.pos_aref + efc.margin, efc.margin - aref, r, pos = fn(efc) - d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref, efc_pos=pos) + aref, r, pos, margin = fn(efc) + d = d.replace( + efc_J=efc.J, efc_D=1 / r, efc_aref=aref, efc_pos=pos, efc_margin=margin + ) d = d.replace(efc_frictionloss=jp.zeros_like(r)) return d diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index 07a57475..b04ec6a1 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -67,7 +67,6 @@ class ConstraintTest(parameterized.TestCase): _assert_eq(d.efc_D, dx.efc_D[order][:d.nefc], 'efc_D') _assert_eq(d.efc_pos, dx.efc_pos[order][:d.nefc], 'efc_pos') - def test_disable_refsafe(self): m = test_util.load_test_file('constraints.xml') @@ -114,6 +113,33 @@ class ConstraintTest(parameterized.TestCase): dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m)) self.assertEqual(dx.efc_J.shape[0], 16) # only equality, joint/tendon limit + def test_margin(self): + """Test margin.""" + m = mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + + + """) + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx = mjx.put_model(m) + dx = mjx.put_data(m, d) + dx = mjx.make_constraint(mx, dx) + + order = test_util.efc_order(m, d, dx) + _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos') + _assert_eq(d.efc_margin, dx.efc_margin[order][: d.nefc], 'efc_margin') + if __name__ == '__main__': absltest.main() diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 4caf0be9..cdd7b2a5 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -286,6 +286,7 @@ def make_data( 'cfrc_ext': (m.nbody, 6, float), 'efc_J': (nefc, m.nv, float), 'efc_pos': (nefc, float), + 'efc_margin': (nefc, float), 'efc_frictionloss': (nefc, float), 'efc_D': (nefc, float), 'efc_aref': (nefc, float), @@ -521,6 +522,7 @@ def put_data( for fname in ( 'efc_J', 'efc_pos', + 'efc_margin', 'efc_frictionloss', 'efc_D', 'efc_aref', diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index eaf54fd4..8b2bd305 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -1201,6 +1201,7 @@ class Data(PyTreeNode): efc_type: constraint type (nefc,) efc_J: constraint Jacobian (nefc, nv) efc_pos: constraint position (equality, contact) (nefc,) + efc_margin: inclusion margin (contact) (nefc,) efc_frictionloss: frictionloss (friction) (nefc,) efc_D: constraint mass (nefc,) efc_aref: reference pseudo-acceleration (nefc,) @@ -1322,6 +1323,7 @@ class Data(PyTreeNode): efc_type: jax.Array efc_J: jax.Array # pylint:disable=invalid-name efc_pos: jax.Array + efc_margin: jax.Array efc_frictionloss: jax.Array efc_D: jax.Array # pylint:disable=invalid-name # dynamically sized - position & velocity dependent: From 8b03daa09b0f1404cb2b968d83ad6a2723b29db3 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Thu, 22 Aug 2024 05:16:02 -0700 Subject: [PATCH 19/26] Add `mjtSameFrame` enum to specify frame alignment of bodies with their children. PiperOrigin-RevId: 666303674 Change-Id: I53dc56e1ca52afa280ae5abaf4158edc465fb1e0 --- doc/APIreference/APItypes.rst | 10 ++ doc/APIreference/functions.rst | 2 +- doc/changelog.rst | 25 +++-- doc/includes/references.h | 11 +- include/mujoco/mjmodel.h | 13 ++- include/mujoco/mujoco.h | 2 +- introspect/enums.py | 10 ++ introspect/functions.py | 2 +- introspect/structs.py | 6 +- src/engine/engine_support.c | 37 +++---- src/user/user_model.cc | 158 ++++++++++++++++----------- src/user/user_model.h | 1 - unity/Runtime/Bindings/MjBindings.cs | 5 + 13 files changed, 173 insertions(+), 109 deletions(-) diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst index 5f58d803..52f2b317 100644 --- a/doc/APIreference/APItypes.rst +++ b/doc/APIreference/APItypes.rst @@ -346,6 +346,16 @@ These are the possible sensor data types, used in ``mjData.sensor_datatype``. .. mujoco-include:: mjtDataType +.. _mjtSameFrame: + +mjtSameFrame +~~~~~~~~~~~~ + +Types of frame alignment of elements with their parent bodies. Used as shortcuts during :ref:`mj_kinematics` in the +last argument to :ref:`mj_local2global`. + +.. mujoco-include:: mjtSameFrame + .. _tyDataEnums: diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index f93a31b7..96ade764 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -551,7 +551,7 @@ mj_local2Global .. mujoco-include:: mj_local2Global -Map from body local to global Cartesian coordinates. +Map from body local to global Cartesian coordinates, sameframe takes values from mjtSameFrame. .. _mj_getTotalmass: diff --git a/doc/changelog.rst b/doc/changelog.rst index 195d20cd..827ecfce 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -12,31 +12,34 @@ General early stages of testing. 2. Added :ref:`mjSpec` option for creating a texture from a buffer. 3. :ref:`shellinertia ` is now supported by all geom types. -4. Added support for :ref:`attaching` keyframes. +4. When :ref:`attaching` sub-models, :ref:`keyframes` will now be correctly merged into the + parent model, but only on the first attachment. +5. Added the :ref:`mjtSameFrame` enum which contains the possible frame alignments of bodies and their children. These + alignments are used as shortcuts in :ref:`mj_kinematics`. MJX ^^^ -5. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`). -6. Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, +6. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`). +7. Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``. -7. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. -8. Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``. -9. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. -10. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. -11. Added support for :ref:`implicitfast integration` for all cases except +8. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. +9. Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``. +10. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. +11. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. +12. Added support for :ref:`implicitfast integration` for all cases except :doc:`fluid drag `. Bug fixes ^^^^^^^^^ -12. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, +13. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, contribution by :github:user:`michael-ahn`). -13. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit +14. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit integrators, wrong derivatives would be computed. Python bindings ^^^^^^^^^^^^^^^ -14. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). +15. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). Version 3.2.2 (Aug 8, 2024) diff --git a/doc/includes/references.h b/doc/includes/references.h index 01942d4e..0ca6cacd 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -684,6 +684,11 @@ typedef enum mjtDataType_ { // data type for sensors mjDATATYPE_AXIS, // 3D unit vector mjDATATYPE_QUATERNION // unit quaternion } mjtDataType; +typedef enum mjtSameFrame_ { // frame alignment of bodies with their children + mjSAMEFRAME_NONE = 0, // no alignment + mjSAMEFRAME_BODY, // frame is same as body frame + mjSAMEFRAME_INERTIA, // frame is same as inertial frame +} mjtSameFrame; typedef enum mjtLRMode_ { // mode for actuator length range computation mjLRMODE_NONE = 0, // do not process any actuators mjLRMODE_MUSCLE, // process muscle actuators @@ -978,7 +983,7 @@ struct mjModel_ { int* body_geomnum; // number of geoms (nbody x 1) int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1) mjtByte* body_simple; // 1: diag M; 2: diag M, sliders only (nbody x 1) - mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1) + mjtByte* body_sameframe; // same frame as inertia (mjtSameframe) (nbody x 1) mjtNum* body_pos; // position offset rel. to parent body (nbody x 3) mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4) mjtNum* body_ipos; // local position of center of mass (nbody x 3) @@ -1047,7 +1052,7 @@ struct mjModel_ { int* geom_group; // group for visibility (ngeom x 1) int* geom_priority; // geom contact priority (ngeom x 1) int* geom_plugin; // plugin instance id; -1: not in use (ngeom x 1) - mjtByte* geom_sameframe; // same as body frame (1) or iframe (2) (ngeom x 1) + mjtByte* geom_sameframe; // same frame as body (mjtSameframe) (ngeom x 1) mjtNum* geom_solmix; // mixing coef for solref/imp in geom pair (ngeom x 1) mjtNum* geom_solref; // constraint solver reference: contact (ngeom x mjNREF) mjtNum* geom_solimp; // constraint solver impedance: contact (ngeom x mjNIMP) @@ -1068,7 +1073,7 @@ struct mjModel_ { int* site_bodyid; // id of site's body (nsite x 1) int* site_matid; // material id for rendering; -1: none (nsite x 1) int* site_group; // group for visibility (nsite x 1) - mjtByte* site_sameframe; // same as body frame (1) or iframe (2) (nsite x 1) + mjtByte* site_sameframe; // same frame as body (mjtSameframe) (nsite x 1) mjtNum* site_size; // geom size for rendering (nsite x 3) mjtNum* site_pos; // local position offset rel. to body (nsite x 3) mjtNum* site_quat; // local orientation offset rel. to body (nsite x 4) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 37e346b7..4fd691c9 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -373,6 +373,13 @@ typedef enum mjtDataType_ { // data type for sensors } mjtDataType; +typedef enum mjtSameFrame_ { // frame alignment of bodies with their children + mjSAMEFRAME_NONE = 0, // no alignment + mjSAMEFRAME_BODY, // frame is same as body frame + mjSAMEFRAME_INERTIA, // frame is same as inertial frame +} mjtSameFrame; + + typedef enum mjtLRMode_ { // mode for actuator length range computation mjLRMODE_NONE = 0, // do not process any actuators mjLRMODE_MUSCLE, // process muscle actuators @@ -692,7 +699,7 @@ struct mjModel_ { int* body_geomnum; // number of geoms (nbody x 1) int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1) mjtByte* body_simple; // 1: diag M; 2: diag M, sliders only (nbody x 1) - mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1) + mjtByte* body_sameframe; // same frame as inertia (mjtSameframe) (nbody x 1) mjtNum* body_pos; // position offset rel. to parent body (nbody x 3) mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4) mjtNum* body_ipos; // local position of center of mass (nbody x 3) @@ -761,7 +768,7 @@ struct mjModel_ { int* geom_group; // group for visibility (ngeom x 1) int* geom_priority; // geom contact priority (ngeom x 1) int* geom_plugin; // plugin instance id; -1: not in use (ngeom x 1) - mjtByte* geom_sameframe; // same as body frame (1) or iframe (2) (ngeom x 1) + mjtByte* geom_sameframe; // same frame as body (mjtSameframe) (ngeom x 1) mjtNum* geom_solmix; // mixing coef for solref/imp in geom pair (ngeom x 1) mjtNum* geom_solref; // constraint solver reference: contact (ngeom x mjNREF) mjtNum* geom_solimp; // constraint solver impedance: contact (ngeom x mjNIMP) @@ -782,7 +789,7 @@ struct mjModel_ { int* site_bodyid; // id of site's body (nsite x 1) int* site_matid; // material id for rendering; -1: none (nsite x 1) int* site_group; // group for visibility (nsite x 1) - mjtByte* site_sameframe; // same as body frame (1) or iframe (2) (nsite x 1) + mjtByte* site_sameframe; // same frame as body (mjtSameframe) (nsite x 1) mjtNum* site_size; // geom size for rendering (nsite x 3) mjtNum* site_pos; // local position offset rel. to body (nsite x 3) mjtNum* site_quat; // local orientation offset rel. to body (nsite x 4) diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 86897462..26014657 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -507,7 +507,7 @@ MJAPI void mj_integratePos(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, m // Normalize all quaternions in qpos-type vector. MJAPI void mj_normalizeQuat(const mjModel* m, mjtNum* qpos); -// Map from body local to global Cartesian coordinates. +// Map from body local to global Cartesian coordinates, sameframe takes values from mjtSameFrame. MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], const mjtNum pos[3], const mjtNum quat[4], int body, mjtByte sameframe); diff --git a/introspect/enums.py b/introspect/enums.py index 22e88ae3..3c36b5cf 100644 --- a/introspect/enums.py +++ b/introspect/enums.py @@ -388,6 +388,16 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjDATATYPE_QUATERNION', 3), ]), )), + ('mjtSameFrame', + EnumDecl( + name='mjtSameFrame', + declname='enum mjtSameFrame_', + values=dict([ + ('mjSAMEFRAME_NONE', 0), + ('mjSAMEFRAME_BODY', 1), + ('mjSAMEFRAME_INERTIA', 2), + ]), + )), ('mjtLRMode', EnumDecl( name='mjtLRMode', diff --git a/introspect/functions.py b/introspect/functions.py index 4e3073e8..d787bf1c 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -3168,7 +3168,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ type=ValueType(name='mjtByte'), ), ), - doc='Map from body local to global Cartesian coordinates.', + doc='Map from body local to global Cartesian coordinates, sameframe takes values from mjtSameFrame.', # pylint: disable=line-too-long )), ('mj_getTotalmass', FunctionDecl( diff --git a/introspect/structs.py b/introspect/structs.py index 2ed6deef..10cccc46 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -1343,7 +1343,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='mjtByte'), ), - doc='inertial frame is same as body frame (nbody x 1)', + doc='same frame as inertia (mjtSameframe) (nbody x 1)', ), StructFieldDecl( name='body_pos', @@ -1770,7 +1770,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='mjtByte'), ), - doc='same as body frame (1) or iframe (2) (ngeom x 1)', + doc='same frame as body (mjtSameframe) (ngeom x 1)', ), StructFieldDecl( name='geom_solmix', @@ -1903,7 +1903,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=PointerType( inner_type=ValueType(name='mjtByte'), ), - doc='same as body frame (1) or iframe (2) (nsite x 1)', + doc='same frame as body (mjtSameframe) (nsite x 1)', ), StructFieldDecl( name='site_size', diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 07bf9ed5..bd6699d5 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -1614,43 +1614,38 @@ void mj_normalizeQuat(const mjModel* m, mjtNum* qpos) { void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], const mjtNum pos[3], const mjtNum quat[4], int body, mjtByte sameframe) { - mjtNum tmp[4]; + mjtSameFrame sf = sameframe; // position if (xpos && pos) { - // compute - if (sameframe == 0) { + switch (sf) { + case mjSAMEFRAME_NONE: mju_mulMatVec3(xpos, d->xmat+9*body, pos); mju_addTo3(xpos, d->xpos+3*body); - } - - // copy body position - else if (sameframe == 1) { + break; + case mjSAMEFRAME_BODY: mju_copy3(xpos, d->xpos+3*body); - } - - // copy inertial body position - else { + break; + case mjSAMEFRAME_INERTIA: mju_copy3(xpos, d->xipos+3*body); + break; } } // orientation if (xmat && quat) { - // compute - if (sameframe == 0) { + mjtNum tmp[4]; + switch (sf) { + case mjSAMEFRAME_NONE: mju_mulQuat(tmp, d->xquat+4*body, quat); mju_quat2Mat(xmat, tmp); - } - - // copy body orientation - else if (sameframe == 1) { + break; + case mjSAMEFRAME_BODY: mju_copy(xmat, d->xmat+9*body, 9); - } - - // copy inertial body orientation - else { + break; + case mjSAMEFRAME_INERTIA: mju_copy(xmat, d->ximat+9*body, 9); + break; } } } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index cbd37927..140156ba 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -19,7 +19,6 @@ #include #include #include -#include #include #include #include @@ -56,6 +55,65 @@ namespace mju = ::mujoco::util; using std::string; using std::vector; constexpr int kMaxCompilerThreads = 16; + + + +//---------------------------------- LOCAL UTILITY FUNCTIONS --------------------------------------- + +constexpr double kFrameEps = 1e-6; // difference below which frames are considered equal + +// return true if two 3-vectors are element-wise less than kFrameEps apart +template +bool IsSameVec(const T pos1[3], const T pos2[3]) { + static_assert(std::is_floating_point_v); + return std::abs(pos1[0] - pos2[0]) < kFrameEps && + std::abs(pos1[1] - pos2[1]) < kFrameEps && + std::abs(pos1[2] - pos2[2]) < kFrameEps; +} + +// return true if two quaternions are element-wise less than kFrameEps apart, including double-cover +template +bool IsSameQuat(const T quat1[4], const T quat2[4]) { + static_assert(std::is_floating_point_v); + bool same_quat_minus = std::abs(quat1[0] - quat2[0]) < kFrameEps && + std::abs(quat1[1] - quat2[1]) < kFrameEps && + std::abs(quat1[2] - quat2[2]) < kFrameEps && + std::abs(quat1[3] - quat2[3]) < kFrameEps; + + bool same_quat_plus = std::abs(quat1[0] + quat2[0]) < kFrameEps && + std::abs(quat1[1] + quat2[1]) < kFrameEps && + std::abs(quat1[2] + quat2[2]) < kFrameEps && + std::abs(quat1[3] + quat2[3]) < kFrameEps; + + return same_quat_minus || same_quat_plus; +} + + +// compare two poses +template +bool IsSamePose(const T pos1[3], const T pos2[3], const T quat1[4], const T quat2[4]) { + // check position if given + if (pos1 && pos2 && !IsSameVec(pos1, pos2)) { + return false; + } + + // check orientation if given + if (quat1 && quat2 && !IsSameQuat(quat1, quat2)) { + return false; + } + + return true; +} + +// detect null pose +template +bool IsNullPose(const T pos[3], const T quat[4]) { + T zero[3] = {0, 0, 0}; + T qunit[4] = {1, 0, 0, 0}; + return IsSamePose(pos, zero, quat, qunit); +} + + } // namespace //---------------------------------- CONSTRUCTOR AND DESTRUCTOR ------------------------------------ @@ -429,7 +487,7 @@ void mjCModel::DeleteElement(mjsElement* el) { switch (el->elemtype) { case mjOBJ_BODY: - throw mjCError(NULL, "bodies cannot be deleted, use detach instead"); + throw mjCError(nullptr, "bodies cannot be deleted, use detach instead"); break; case mjOBJ_GEOM: @@ -825,12 +883,12 @@ static mjsElement* GetNext(std::vector& list, mjsElement* child) { mjsElement* mjCModel::NextObject(mjsElement* object, mjtObj type) { if (type == mjOBJ_UNKNOWN) { if (!object) { - throw mjCError(NULL, "type must be specified if no element is given"); + throw mjCError(nullptr, "type must be specified if no element is given"); } else { type = object->elemtype; } } else if (object && object->elemtype != type) { - throw mjCError(NULL, "element is not of requested type"); + throw mjCError(nullptr, "element is not of requested type"); } switch (type) { @@ -1022,29 +1080,6 @@ mjSpec* mjCModel::FindSpec(std::string name) const { -// detect null pose -bool mjCModel::IsNullPose(const mjtNum* pos, const mjtNum* quat) const { - bool result = true; - - // check position if given - if (pos) { - if (pos[0] || pos[1] || pos[2]) { - result = false; - } - } - - // check orientation if given - if (quat) { - if (quat[0]!=1 || quat[1] || quat[2] || quat[3]) { - result = false; - } - } - - return result; -} - - - //------------------------------- COMPILER PHASES -------------------------------------------------- // make lists of objects in tree: bodies, geoms, joints, sites, cameras, lights @@ -1640,12 +1675,12 @@ void* LRfunc(void* arg) { for (int i=larg->start; istart+larg->num; i++) { if (im->nu) { if (!mj_setLengthRange(larg->m, larg->data, i, larg->LRopt, larg->error, larg->error_sz)) { - return NULL; + return nullptr; } } } - return NULL; + return nullptr; } @@ -1977,16 +2012,22 @@ void mjCModel::CopyTree(mjModel* m) { pb->lastdof = par->lastdof; // set sameframe - m->body_sameframe[i] = IsNullPose(m->body_ipos+3*i, m->body_iquat+4*i); + mjtSameFrame sameframe; + if (IsNullPose(m->body_ipos+3*i, m->body_iquat+4*i)) { + sameframe = mjSAMEFRAME_BODY; + } else { + sameframe = mjSAMEFRAME_NONE; + } + m->body_sameframe[i] = sameframe; // init simple: sameframe, and (self-root, or parent is fixed child of world) - int j = m->body_parentid[i]; - m->body_simple[i] = (m->body_sameframe[i] && + int parentid = m->body_parentid[i]; + m->body_simple[i] = (sameframe == mjSAMEFRAME_BODY && (m->body_rootid[i]==i || - (m->body_parentid[j]==0 && - m->body_dofnum[j]==0))); + (m->body_parentid[parentid]==0 && + m->body_dofnum[parentid]==0))); - // parent is not simple (unless world) + // a parent body is never simple (unless world) if (m->body_parentid[i]>0) { m->body_simple[m->body_parentid[i]] = 0; } @@ -2020,12 +2061,11 @@ void mjCModel::CopyTree(mjModel* m) { mjuu_copyvec(m->jnt_user+nuser_jnt*jid, pj->get_userdata().data(), nuser_jnt); // not simple if: rotation already found, or pos not zero, or mis-aligned axis - if (rotfound || - !IsNullPose(m->jnt_pos+3*jid, NULL) || - ((pj->type==mjJNT_HINGE || pj->type==mjJNT_SLIDE) && - ((std::abs(pj->axis[0])>mjEPS) + - (std::abs(pj->axis[1])>mjEPS) + - (std::abs(pj->axis[2])>mjEPS)) > 1)) { + bool axis_aligned = ((std::abs(pj->axis[0]) > mjEPS) + + (std::abs(pj->axis[1]) > mjEPS) + + (std::abs(pj->axis[2]) > mjEPS)) == 1; + if (rotfound || !IsNullPose(m->jnt_pos+3*jid, static_cast(nullptr)) || + ((pj->type == mjJNT_HINGE || pj->type == mjJNT_SLIDE) && !axis_aligned)) { m->body_simple[i] = 0; } @@ -2133,18 +2173,13 @@ void mjCModel::CopyTree(mjModel* m) { // determine sameframe if (IsNullPose(m->geom_pos+3*gid, m->geom_quat+4*gid)) { - m->geom_sameframe[gid] = 1; - } else if (pg->pos[0]==pb->ipos[0] && - pg->pos[1]==pb->ipos[1] && - pg->pos[2]==pb->ipos[2] && - pg->quat[0]==pb->iquat[0] && - pg->quat[1]==pb->iquat[1] && - pg->quat[2]==pb->iquat[2] && - pg->quat[3]==pb->iquat[3]) { - m->geom_sameframe[gid] = 2; + sameframe = mjSAMEFRAME_BODY; + } else if (IsSamePose(pg->pos, pb->ipos, pg->quat, pb->iquat)) { + sameframe = mjSAMEFRAME_INERTIA; } else { - m->geom_sameframe[gid] = 0; + sameframe = mjSAMEFRAME_NONE; } + m->geom_sameframe[gid] = sameframe; // compute rbound m->geom_rbound[gid] = (mjtNum)pg->GetRBound(); @@ -2169,18 +2204,13 @@ void mjCModel::CopyTree(mjModel* m) { // determine sameframe if (IsNullPose(m->site_pos+3*sid, m->site_quat+4*sid)) { - m->site_sameframe[sid] = 1; - } else if (ps->pos[0]==pb->ipos[0] && - ps->pos[1]==pb->ipos[1] && - ps->pos[2]==pb->ipos[2] && - ps->quat[0]==pb->iquat[0] && - ps->quat[1]==pb->iquat[1] && - ps->quat[2]==pb->iquat[2] && - ps->quat[3]==pb->iquat[3]) { - m->site_sameframe[sid] = 2; + sameframe = mjSAMEFRAME_BODY; + } else if (IsSamePose(ps->pos, pb->ipos, ps->quat, pb->iquat)) { + sameframe = mjSAMEFRAME_INERTIA; } else { - m->site_sameframe[sid] = 0; + sameframe = mjSAMEFRAME_NONE; } + m->site_sameframe[sid] = sameframe; } // loop over cameras for this body @@ -2887,7 +2917,7 @@ void mjCModel::SaveState(const std::string& state_name, const T* qpos, const T* const T* ctrl, const T* mpos, const T* mquat) { for (auto joint : joints_) { if (joint->qposadr_ == -1 || joint->dofadr_ == -1) { - throw mjCError(NULL, "SaveState: joint %s has no address", joint->name.c_str()); + throw mjCError(nullptr, "SaveState: joint %s has no address", joint->name.c_str()); } if (qpos) mjuu_copyvec(joint->qpos(state_name), qpos + joint->qposadr_, joint->nq()); if (qvel) mjuu_copyvec(joint->qvel(state_name), qvel + joint->dofadr_, joint->nv()); @@ -3170,7 +3200,7 @@ void mjCModel::FuseStatic(void) { mjuu_copyvec(par->fullinertia, toti, 6); const char* err1 = mjuu_fullInertia(par->iquat, par->inertia, par->fullinertia); if (err1) { - throw mjCError(NULL, "error '%s' in fusing static body inertias", err1); + throw mjCError(nullptr, "error '%s' in fusing static body inertias", err1); } } } @@ -3348,7 +3378,7 @@ static void processlist(mjListKeyMap& ids, vector& list, auto adjacent = std::adjacent_find(allnames.begin(), allnames.end()); if (adjacent != allnames.end()) { string msg = "repeated name '" + *adjacent + "' in " + mju_type2Str(type); - throw mjCError(NULL, "%s", msg.c_str()); + throw mjCError(nullptr, "%s", msg.c_str()); } } } diff --git a/src/user/user_model.h b/src/user/user_model.h index f4ccf839..4529b8f7 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -232,7 +232,6 @@ class mjCModel : public mjCModel_, private mjSpec { mjCBody* FindBody(mjCBody* body, std::string name); // find body given name mjCFrame* FindFrame(mjCBody* body, std::string name) const; // find frame given name mjSpec* FindSpec(std::string name) const; // find spec given name - bool IsNullPose(const mjtNum* pos, const mjtNum* quat) const; // detect null pose void SetActivePlugins(const std::vector>&& active_plugins) { active_plugins_ = std::move(active_plugins); } diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 3f26168a..0d62e552 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -392,6 +392,11 @@ public enum mjtDataType : int{ mjDATATYPE_AXIS = 2, mjDATATYPE_QUATERNION = 3, } +public enum mjtSameFrame : int{ + mjSAMEFRAME_NONE = 0, + mjSAMEFRAME_BODY = 1, + mjSAMEFRAME_INERTIA = 2, +} public enum mjtLRMode : int{ mjLRMODE_NONE = 0, mjLRMODE_MUSCLE = 1, From 2efbd317628bb45ec6cc12104224a9e51feffd10 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Thu, 22 Aug 2024 06:07:30 -0700 Subject: [PATCH 20/26] Add `mjSAMEFRAME_BODYROT` and `mjSAMEFRAME_INERTIAROT` frame alignment shortcuts. PiperOrigin-RevId: 666318799 Change-Id: I0af1824d2bd9eaca7b48a9e33e09f873dd33b22d --- doc/includes/references.h | 2 + include/mujoco/mjmodel.h | 2 + introspect/enums.py | 2 + src/engine/engine_support.c | 4 ++ src/user/user_model.cc | 15 +++++++- test/user/user_model_test.cc | 57 ++++++++++++++++++++++++++++ unity/Runtime/Bindings/MjBindings.cs | 2 + 7 files changed, 83 insertions(+), 1 deletion(-) diff --git a/doc/includes/references.h b/doc/includes/references.h index 0ca6cacd..87010fe4 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -688,6 +688,8 @@ typedef enum mjtSameFrame_ { // frame alignment of bodies with their childr mjSAMEFRAME_NONE = 0, // no alignment mjSAMEFRAME_BODY, // frame is same as body frame mjSAMEFRAME_INERTIA, // frame is same as inertial frame + mjSAMEFRAME_BODYROT, // frame orientation is same as body orientation + mjSAMEFRAME_INERTIAROT // frame orientation is same as inertia orientation } mjtSameFrame; typedef enum mjtLRMode_ { // mode for actuator length range computation mjLRMODE_NONE = 0, // do not process any actuators diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 4fd691c9..16e18ea8 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -377,6 +377,8 @@ typedef enum mjtSameFrame_ { // frame alignment of bodies with their childr mjSAMEFRAME_NONE = 0, // no alignment mjSAMEFRAME_BODY, // frame is same as body frame mjSAMEFRAME_INERTIA, // frame is same as inertial frame + mjSAMEFRAME_BODYROT, // frame orientation is same as body orientation + mjSAMEFRAME_INERTIAROT // frame orientation is same as inertia orientation } mjtSameFrame; diff --git a/introspect/enums.py b/introspect/enums.py index 3c36b5cf..9e24b3b9 100644 --- a/introspect/enums.py +++ b/introspect/enums.py @@ -396,6 +396,8 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjSAMEFRAME_NONE', 0), ('mjSAMEFRAME_BODY', 1), ('mjSAMEFRAME_INERTIA', 2), + ('mjSAMEFRAME_BODYROT', 3), + ('mjSAMEFRAME_INERTIAROT', 4), ]), )), ('mjtLRMode', diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index bd6699d5..35510f8d 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -1620,6 +1620,8 @@ void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], if (xpos && pos) { switch (sf) { case mjSAMEFRAME_NONE: + case mjSAMEFRAME_BODYROT: + case mjSAMEFRAME_INERTIAROT: mju_mulMatVec3(xpos, d->xmat+9*body, pos); mju_addTo3(xpos, d->xpos+3*body); break; @@ -1641,9 +1643,11 @@ void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9], mju_quat2Mat(xmat, tmp); break; case mjSAMEFRAME_BODY: + case mjSAMEFRAME_BODYROT: mju_copy(xmat, d->xmat+9*body, 9); break; case mjSAMEFRAME_INERTIA: + case mjSAMEFRAME_INERTIAROT: mju_copy(xmat, d->ximat+9*body, 9); break; } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 140156ba..646ecfe7 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -2013,8 +2013,11 @@ void mjCModel::CopyTree(mjModel* m) { // set sameframe mjtSameFrame sameframe; + mjtNum* nullnum = static_cast(nullptr); if (IsNullPose(m->body_ipos+3*i, m->body_iquat+4*i)) { sameframe = mjSAMEFRAME_BODY; + } else if (IsNullPose(nullnum, m->body_iquat+4*i)) { + sameframe = mjSAMEFRAME_BODYROT; } else { sameframe = mjSAMEFRAME_NONE; } @@ -2064,7 +2067,7 @@ void mjCModel::CopyTree(mjModel* m) { bool axis_aligned = ((std::abs(pj->axis[0]) > mjEPS) + (std::abs(pj->axis[1]) > mjEPS) + (std::abs(pj->axis[2]) > mjEPS)) == 1; - if (rotfound || !IsNullPose(m->jnt_pos+3*jid, static_cast(nullptr)) || + if (rotfound || !IsNullPose(m->jnt_pos+3*jid, nullnum) || ((pj->type == mjJNT_HINGE || pj->type == mjJNT_SLIDE) && !axis_aligned)) { m->body_simple[i] = 0; } @@ -2172,10 +2175,15 @@ void mjCModel::CopyTree(mjModel* m) { mjuu_copyvec(m->geom_rgba+4*gid, pg->rgba, 4); // determine sameframe + double* nulldouble = static_cast(nullptr); if (IsNullPose(m->geom_pos+3*gid, m->geom_quat+4*gid)) { sameframe = mjSAMEFRAME_BODY; + } else if (IsNullPose(nullnum, m->geom_quat+4*gid)) { + sameframe = mjSAMEFRAME_BODYROT; } else if (IsSamePose(pg->pos, pb->ipos, pg->quat, pb->iquat)) { sameframe = mjSAMEFRAME_INERTIA; + } else if (IsSamePose(nulldouble, nulldouble, pg->quat, pb->iquat)) { + sameframe = mjSAMEFRAME_INERTIAROT; } else { sameframe = mjSAMEFRAME_NONE; } @@ -2203,10 +2211,15 @@ void mjCModel::CopyTree(mjModel* m) { mjuu_copyvec(m->site_rgba+4*sid, ps->rgba, 4); // determine sameframe + double* nulldouble = static_cast(nullptr); if (IsNullPose(m->site_pos+3*sid, m->site_quat+4*sid)) { sameframe = mjSAMEFRAME_BODY; + } else if (IsNullPose(nullnum, m->site_quat+4*sid)) { + sameframe = mjSAMEFRAME_BODYROT; } else if (IsSamePose(ps->pos, pb->ipos, ps->quat, pb->iquat)) { sameframe = mjSAMEFRAME_INERTIA; + } else if (IsSamePose(nulldouble, nulldouble, ps->quat, pb->iquat)) { + sameframe = mjSAMEFRAME_INERTIAROT; } else { sameframe = mjSAMEFRAME_NONE; } diff --git a/test/user/user_model_test.cc b/test/user/user_model_test.cc index 8b4f2cdb..6e6f1bbc 100644 --- a/test/user/user_model_test.cc +++ b/test/user/user_model_test.cc @@ -36,12 +36,17 @@ using ::testing::ElementsAre; using ::testing::HasSubstr; using ::testing::IsNull; using ::testing::NotNull; +using ::testing::Pointwise; static std::vector GetRow(const mjtNum* array, int ncolumn, int row) { return std::vector(array + ncolumn * row, array + ncolumn * (row + 1)); } +std::vector AsVector(const mjtNum* array, int n) { + return std::vector(array, array + n); +} + // ----------------------------- test mjCModel -------------------------------- using UserCModelTest = MujocoTest; @@ -64,6 +69,58 @@ TEST_F(UserCModelTest, RepeatedNames) { EXPECT_THAT(error.data(), HasSubstr("repeated name 'geom1' in geom")); } +TEST_F(UserCModelTest, SameFrame) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + )"; + + std::array error; + mjModel* model = LoadModelFromString(xml, error.data(), error.size()); + ASSERT_THAT(model, NotNull()) << error.data(); + EXPECT_EQ(model->geom_sameframe[0], mjSAMEFRAME_NONE); + EXPECT_EQ(model->geom_sameframe[1], mjSAMEFRAME_BODY); + EXPECT_EQ(model->geom_sameframe[2], mjSAMEFRAME_INERTIA); + EXPECT_EQ(model->geom_sameframe[3], mjSAMEFRAME_BODYROT); + EXPECT_EQ(model->geom_sameframe[4], mjSAMEFRAME_INERTIAROT); + + // make data, get geom_xpos + mjData* data = mj_makeData(model); + mj_kinematics(model, data); + auto geom_xpos = AsVector(data->geom_xpos, model->ngeom*3); + auto geom_xmat = AsVector(data->geom_xmat, model->ngeom*9); + + // set all geom_sameframe to 0, call kinematics again + for (int i = 0; i < model->ngeom; i++) { + model->geom_sameframe[i] = mjSAMEFRAME_NONE; + } + mj_resetData(model, data); + mj_kinematics(model, data); + auto geom_xpos2 = AsVector(data->geom_xpos, model->ngeom*3); + auto geom_xmat2 = AsVector(data->geom_xmat, model->ngeom*9); + + // expect them to be equal + constexpr double eps = 1e-6; + EXPECT_THAT(geom_xpos, Pointwise(DoubleNear(eps), geom_xpos2)); + EXPECT_THAT(geom_xmat, Pointwise(DoubleNear(eps), geom_xmat2)); + + mj_deleteData(data); + mj_deleteModel(model); +} + + // ------------- test automatic inference of nuser_xxx ------------------------- using UserDataTest = MujocoTest; diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 0d62e552..08ebfe4e 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -396,6 +396,8 @@ public enum mjtSameFrame : int{ mjSAMEFRAME_NONE = 0, mjSAMEFRAME_BODY = 1, mjSAMEFRAME_INERTIA = 2, + mjSAMEFRAME_BODYROT = 3, + mjSAMEFRAME_INERTIAROT = 4, } public enum mjtLRMode : int{ mjLRMODE_NONE = 0, From abe57c234baa5f3776a0fc9aff5f86f221532109 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Thu, 22 Aug 2024 07:09:51 -0700 Subject: [PATCH 21/26] Use bullets instead of numbering in latest changelog PiperOrigin-RevId: 666335980 Change-Id: Ifbeba31e34e9891f4d0791aa4b014d1fa1403398 --- doc/changelog.rst | 48 +++++++++++++++++++++++------------------------ 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 827ecfce..dc0467ba 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -7,39 +7,39 @@ Upcoming version (not yet released) General ^^^^^^^ -1. Added the :ref:`nativeccd` flag. When this flag is enabled, general convex collision - detection is handled natively, as opposed to using `libccd `__. This feature is in - early stages of testing. -2. Added :ref:`mjSpec` option for creating a texture from a buffer. -3. :ref:`shellinertia ` is now supported by all geom types. -4. When :ref:`attaching` sub-models, :ref:`keyframes` will now be correctly merged into the - parent model, but only on the first attachment. -5. Added the :ref:`mjtSameFrame` enum which contains the possible frame alignments of bodies and their children. These - alignments are used as shortcuts in :ref:`mj_kinematics`. +- Added the :ref:`nativeccd` flag. When this flag is enabled, general convex collision + detection is handled natively, as opposed to using `libccd `__. This feature is in + early stages of testing. +- Added :ref:`mjSpec` option for creating a texture from a buffer. +- :ref:`shellinertia ` is now supported by all geom types. +- When :ref:`attaching` sub-models, :ref:`keyframes` will now be correctly merged into the + parent model, but only on the first attachment. +- Added the :ref:`mjtSameFrame` enum which contains the possible frame alignments of bodies and their children. These + alignments are used as shortcuts in :ref:`mj_kinematics`. MJX ^^^ -6. Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`). -7. Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, - ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, - ``SUBTREECOM``, ``CLOCK``. -8. Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. -9. Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``. -10. Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. -11. Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. -12. Added support for :ref:`implicitfast integration` for all cases except - :doc:`fluid drag `. +- Added ``efc_pos`` to ``mjx.Data`` (:github:issue:`1388`). +- Added position-dependent sensors: ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, + ``ACTUATORPOS``, ``BALLQUAT``, ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, + ``SUBTREECOM``, ``CLOCK``. +- Added velocity-dependent sensors: ``JOINTVEL``, ``ACTUATORVEL``, ``BALLANGVEL``. +- Added acceleration/force-dependent sensors: ``ACTUATORFRC``, ``JOINTACTFRC``. +- Changed default policy to avoid placing unused (MuJoCo-only) arrays on device. +- Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``. +- Added support for :ref:`implicitfast integration` for all cases except + :doc:`fluid drag `. Bug fixes ^^^^^^^^^ -13. Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, - contribution by :github:user:`michael-ahn`). -14. Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit - integrators, wrong derivatives would be computed. +- Fixed a performance regression introduced in 3.1.7 in mesh Bounding Volume Hierarchies (:github:issue:`1875`, + contribution by :github:user:`michael-ahn`). +- Fixed a bug wherein, for models that have both muscles and stateless actuators and used one of the implicit + integrators, wrong derivatives would be computed. Python bindings ^^^^^^^^^^^^^^^ -15. Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). +- Added support for engine plugins in :ref:`mjSpec` (:github:issue:`1903`). Version 3.2.2 (Aug 8, 2024) From 088079eff0450e32b98ee743141780ed68307506 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Thu, 22 Aug 2024 12:01:04 -0700 Subject: [PATCH 22/26] Rename mpr_tolerance and mpr_iterations to ccd_tolerance and ccd_iterations. PiperOrigin-RevId: 666438658 Change-Id: I23dbdf939b9ab581f0fca4445f95c5bb36f6b463 --- doc/XMLreference.rst | 12 ++++++------ doc/XMLschema.rst | 4 ++-- doc/changelog.rst | 9 +++++++++ doc/includes/references.h | 4 ++-- include/mujoco/mjmodel.h | 4 ++-- include/mujoco/mjxmacro.h | 4 ++-- introspect/structs.py | 8 ++++---- mjx/mujoco/mjx/_src/types.py | 8 ++++---- simulate/simulate.cc | 8 ++++---- src/engine/engine_collision_convex.c | 10 +++++----- src/engine/engine_io.c | 4 ++-- src/xml/xml_native_reader.cc | 8 ++++---- src/xml/xml_native_writer.cc | 4 ++-- unity/Runtime/Bindings/MjBindings.cs | 4 ++-- unity/Runtime/Components/MjGlobalSettings.cs | 16 ++++++++-------- .../Editor/Components/MjGlobalSettingsTests.cs | 16 ++++++++-------- .../Components/MjcfGenerationContextTests.cs | 4 ++-- 17 files changed, 68 insertions(+), 59 deletions(-) diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index bcd283e8..bf4ec37d 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -444,16 +444,16 @@ adjust it properly through the XML. :at:`noslip_tolerance`: :at-val:`real, "1e-6"` Tolerance threshold used for early termination of the Noslip solver. -.. _option-mpr_iterations: +.. _option-ccd_iterations: -:at:`mpr_iterations`: :at-val:`int, "50"` - Maximum number of iterations of the MPR algorithm used for convex mesh collisions. This rarely needs to be adjusted, +:at:`ccd_iterations`: :at-val:`int, "50"` + Maximum number of iterations of the algorithm used for convex collisions. This rarely needs to be adjusted, except in situations where some geoms have very large aspect ratios. -.. _option-mpr_tolerance: +.. _option-ccd_tolerance: -:at:`mpr_tolerance`: :at-val:`real, "1e-6"` - Tolerance threshold used for early termination of the MPR algorithm. +:at:`ccd_tolerance`: :at-val:`real, "1e-6"` + Tolerance threshold used for early termination of the convex collision algorithm. .. _option-sdf_iterations: diff --git a/doc/XMLschema.rst b/doc/XMLschema.rst index 5adf0c74..2cb92df3 100644 --- a/doc/XMLschema.rst +++ b/doc/XMLschema.rst @@ -12,7 +12,7 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`timestep` | :ref:`apirate` | :ref:`impratio` | :ref:`tolerance` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`ls_tolerance` | :ref:`noslip_tolerance` | :ref:`mpr_tolerance` | :ref:`gravity` | | +| | | | :ref:`ls_tolerance` | :ref:`noslip_tolerance` | :ref:`ccd_tolerance` | :ref:`gravity` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`wind` | :ref:`magnetic` | :ref:`density` | :ref:`viscosity` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | @@ -20,7 +20,7 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`integrator` | :ref:`cone` | :ref:`jacobian` | :ref:`solver` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`iterations` | :ref:`ls_iterations` | :ref:`noslip_iterations` | :ref:`mpr_iterations` | | +| | | | :ref:`iterations` | :ref:`ls_iterations` | :ref:`noslip_iterations` | :ref:`ccd_iterations` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`sdf_iterations` | :ref:`sdf_initpoints` | :ref:`actuatorgroupdisable` | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | diff --git a/doc/changelog.rst b/doc/changelog.rst index dc0467ba..368a6348 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -7,6 +7,15 @@ Upcoming version (not yet released) General ^^^^^^^ + +.. admonition:: Breaking API changes + :class: attention + + - The runtime options ``mpr_tolerance`` and ``mpr_iterations`` were renamed to + :ref:`ccd_tolerance` and :ref:`ccd_iterations`, both in XML and in + the :ref:`mjOption` struct. This is because the new convex collision detection pipeline (see below) does not use + the MPR algorithm. The semantics of these options remain identical. + - Added the :ref:`nativeccd` flag. When this flag is enabled, general convex collision detection is handled natively, as opposed to using `libccd `__. This feature is in early stages of testing. diff --git a/doc/includes/references.h b/doc/includes/references.h index 87010fe4..a70bc244 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -734,7 +734,7 @@ struct mjOption_ { // physics options mjtNum tolerance; // main solver tolerance mjtNum ls_tolerance; // CG/Newton linesearch tolerance mjtNum noslip_tolerance; // noslip solver tolerance - mjtNum mpr_tolerance; // MPR solver tolerance + mjtNum ccd_tolerance; // convex collision solver tolerance // physical constants mjtNum gravity[3]; // gravitational acceleration @@ -757,7 +757,7 @@ struct mjOption_ { // physics options int iterations; // maximum number of main solver iterations int ls_iterations; // maximum number of CG/Newton linesearch iterations int noslip_iterations; // maximum number of noslip solver iterations - int mpr_iterations; // maximum number of MPR solver iterations + int ccd_iterations; // maximum number of convex collision solver iterations int disableflags; // bit flags for disabling standard features int enableflags; // bit flags for enabling optional features int disableactuator; // bit flags for disabling actuators by group id diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 16e18ea8..3c536103 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -438,7 +438,7 @@ struct mjOption_ { // physics options mjtNum tolerance; // main solver tolerance mjtNum ls_tolerance; // CG/Newton linesearch tolerance mjtNum noslip_tolerance; // noslip solver tolerance - mjtNum mpr_tolerance; // MPR solver tolerance + mjtNum ccd_tolerance; // convex collision solver tolerance // physical constants mjtNum gravity[3]; // gravitational acceleration @@ -461,7 +461,7 @@ struct mjOption_ { // physics options int iterations; // maximum number of main solver iterations int ls_iterations; // maximum number of CG/Newton linesearch iterations int noslip_iterations; // maximum number of noslip solver iterations - int mpr_iterations; // maximum number of MPR solver iterations + int ccd_iterations; // maximum number of convex collision solver iterations int disableflags; // bit flags for disabling standard features int enableflags; // bit flags for enabling optional features int disableactuator; // bit flags for disabling actuators by group id diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 85f23af1..37dfeb4a 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -26,7 +26,7 @@ X( mjtNum, tolerance ) \ X( mjtNum, ls_tolerance ) \ X( mjtNum, noslip_tolerance ) \ - X( mjtNum, mpr_tolerance ) \ + X( mjtNum, ccd_tolerance ) \ X( mjtNum, density ) \ X( mjtNum, viscosity ) \ X( mjtNum, o_margin ) \ @@ -40,7 +40,7 @@ X( int, iterations ) \ X( int, ls_iterations ) \ X( int, noslip_iterations ) \ - X( int, mpr_iterations ) \ + X( int, ccd_iterations ) \ X( int, disableflags ) \ X( int, enableflags ) \ X( int, disableactuator ) \ diff --git a/introspect/structs.py b/introspect/structs.py index 10cccc46..fe789863 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -135,9 +135,9 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='noslip solver tolerance', ), StructFieldDecl( - name='mpr_tolerance', + name='ccd_tolerance', type=ValueType(name='mjtNum'), - doc='MPR solver tolerance', + doc='convex collision solver tolerance', ), StructFieldDecl( name='gravity', @@ -238,9 +238,9 @@ STRUCTS: Mapping[str, StructDecl] = dict([ doc='maximum number of noslip solver iterations', ), StructFieldDecl( - name='mpr_iterations', + name='ccd_iterations', type=ValueType(name='int'), - doc='maximum number of MPR solver iterations', + doc='maximum number of convex collision solver iterations', ), StructFieldDecl( name='disableflags', diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 8b2bd305..f913914e 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -362,7 +362,7 @@ class Option(PyTreeNode): tolerance: main solver tolerance ls_tolerance: CG/Newton linesearch tolerance noslip_tolerance: noslip solver tolerance (not used) - mpr_tolerance: MPR solver tolerance (not used) + ccd_tolerance: CCD solver tolerance (not used) gravity: gravitational acceleration (3,) wind: wind (for lift, drag and viscosity) magnetic: global magnetic flux (not used) @@ -384,7 +384,7 @@ class Option(PyTreeNode): iterations: number of main solver iterations ls_iterations: maximum number of CG/Newton linesearch iterations noslip_iterations: maximum number of noslip solver iterations (not used) - mpr_iterations: maximum number of MPR solver iterations (not used) + ccd_iterations: maximum number of CCD solver iterations (not used) disableflags: bit flags for disabling standard features enableflags: bit flags for enabling optional features (not used) disableactuator: bit flags for disabling actuators by group id (not used) @@ -397,7 +397,7 @@ class Option(PyTreeNode): tolerance: jax.Array ls_tolerance: jax.Array noslip_tolerance: jax.Array = _restricted_to('mujoco') - mpr_tolerance: jax.Array = _restricted_to('mujoco') + ccd_tolerance: jax.Array = _restricted_to('mujoco') gravity: jax.Array wind: jax.Array magnetic: jax.Array @@ -415,7 +415,7 @@ class Option(PyTreeNode): iterations: int ls_iterations: int noslip_iterations: int = _restricted_to('mujoco') - mpr_iterations: int = _restricted_to('mujoco') + ccd_iterations: int = _restricted_to('mujoco') disableflags: DisableBit enableflags: int disableactuator: int diff --git a/simulate/simulate.cc b/simulate/simulate.cc index d676129b..877366d5 100644 --- a/simulate/simulate.cc +++ b/simulate/simulate.cc @@ -690,8 +690,8 @@ void MakePhysicsSection(mj::Simulate* sim) { {mjITEM_EDITNUM, "LS Tol", 2, &(opt->ls_tolerance), "1 0 0.1"}, {mjITEM_EDITINT, "Noslip Iter", 2, &(opt->noslip_iterations), "1 0 1000"}, {mjITEM_EDITNUM, "Noslip Tol", 2, &(opt->noslip_tolerance), "1 0 1"}, - {mjITEM_EDITINT, "MPR Iter", 2, &(opt->mpr_iterations), "1 0 1000"}, - {mjITEM_EDITNUM, "MPR Tol", 2, &(opt->mpr_tolerance), "1 0 1"}, + {mjITEM_EDITINT, "CCD Iter", 2, &(opt->ccd_iterations), "1 0 1000"}, + {mjITEM_EDITNUM, "CCD Tol", 2, &(opt->ccd_tolerance), "1 0 1"}, {mjITEM_EDITNUM, "API Rate", 2, &(opt->apirate), "1 0 1000"}, {mjITEM_EDITINT, "SDF Iter", 2, &(opt->sdf_iterations), "1 1 20"}, {mjITEM_EDITINT, "SDF Init", 2, &(opt->sdf_initpoints), "1 1 100"}, @@ -1886,7 +1886,7 @@ void Simulate::Sync() { X(impratio); X(tolerance); X(noslip_tolerance); - X(mpr_tolerance); + X(ccd_tolerance); X(gravity); X(wind); X(magnetic); @@ -1902,7 +1902,7 @@ void Simulate::Sync() { X(solver); X(iterations); X(noslip_iterations); - X(mpr_iterations); + X(ccd_iterations); X(disableflags); X(enableflags); X(disableactuator); diff --git a/src/engine/engine_collision_convex.c b/src/engine/engine_collision_convex.c index ceba996b..39535d28 100644 --- a/src/engine/engine_collision_convex.c +++ b/src/engine/engine_collision_convex.c @@ -291,9 +291,9 @@ void mjc_support(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]) { // initialize CCD structure static void mjc_initCCD(ccd_t* ccd, const mjModel* m) { CCD_INIT(ccd); - ccd->mpr_tolerance = m->opt.mpr_tolerance; - ccd->epa_tolerance = m->opt.mpr_tolerance; // use MPR tolerance for EPA - ccd->max_iterations = m->opt.mpr_iterations; + ccd->mpr_tolerance = m->opt.ccd_tolerance; + ccd->epa_tolerance = m->opt.ccd_tolerance; // use MPR tolerance for EPA + ccd->max_iterations = m->opt.ccd_iterations; } @@ -1237,8 +1237,8 @@ int mjc_HFieldElem(const mjModel* m, const mjData* d, mjContact* con, ccd.support2 = mjccd_support; // set ccd parameters - ccd.max_iterations = m->opt.mpr_iterations; - ccd.mpr_tolerance = m->opt.mpr_tolerance; + ccd.max_iterations = m->opt.ccd_iterations; + ccd.mpr_tolerance = m->opt.ccd_tolerance; // compute real-valued grid step, and triangulation direction dx = (2.0*hsize[0]) / (ncol-1); diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 101b7774..07a72c7c 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -124,7 +124,7 @@ void mj_defaultOption(mjOption* opt) { opt->tolerance = 1e-8; opt->ls_tolerance = 0.01; opt->noslip_tolerance = 1e-6; - opt->mpr_tolerance = 1e-6; + opt->ccd_tolerance = 1e-6; // physical constants opt->gravity[0] = 0; @@ -156,7 +156,7 @@ void mj_defaultOption(mjOption* opt) { opt->iterations = 100; opt->ls_iterations = 50; opt->noslip_iterations = 0; - opt->mpr_iterations = 50; + opt->ccd_iterations = 50; opt->disableflags = 0; opt->enableflags = 0; opt->disableactuator = 0; diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc index 85a5f5d2..a6097ceb 100644 --- a/src/xml/xml_native_reader.cc +++ b/src/xml/xml_native_reader.cc @@ -110,10 +110,10 @@ const char* MJCF[nMJCF][mjXATTRNUM] = { {"option", "*", "27", "timestep", "apirate", "impratio", "tolerance", "ls_tolerance", "noslip_tolerance", - "mpr_tolerance", "gravity", "wind", "magnetic", "density", "viscosity", + "ccd_tolerance", "gravity", "wind", "magnetic", "density", "viscosity", "o_margin", "o_solref", "o_solimp", "o_friction", "integrator", "cone", "jacobian", - "solver", "iterations", "ls_iterations", "noslip_iterations", "mpr_iterations", + "solver", "iterations", "ls_iterations", "noslip_iterations", "ccd_iterations", "sdf_iterations", "sdf_initpoints", "actuatorgroupdisable"}, {"<"}, {"flag", "?", "23", "constraint", "equality", "frictionloss", "limit", "contact", @@ -1061,7 +1061,7 @@ void mjXReader::Option(XMLElement* section, mjOption* opt) { ReadAttr(section, "tolerance", 1, &opt->tolerance, text); ReadAttr(section, "ls_tolerance", 1, &opt->ls_tolerance, text); ReadAttr(section, "noslip_tolerance", 1, &opt->noslip_tolerance, text); - ReadAttr(section, "mpr_tolerance", 1, &opt->mpr_tolerance, text); + ReadAttr(section, "ccd_tolerance", 1, &opt->ccd_tolerance, text); ReadAttr(section, "gravity", 3, opt->gravity, text); ReadAttr(section, "wind", 3, opt->wind, text); ReadAttr(section, "magnetic", 3, opt->magnetic, text); @@ -1080,7 +1080,7 @@ void mjXReader::Option(XMLElement* section, mjOption* opt) { ReadAttrInt(section, "iterations", &opt->iterations); ReadAttrInt(section, "ls_iterations", &opt->ls_iterations); ReadAttrInt(section, "noslip_iterations", &opt->noslip_iterations); - ReadAttrInt(section, "mpr_iterations", &opt->mpr_iterations); + ReadAttrInt(section, "ccd_iterations", &opt->ccd_iterations); ReadAttrInt(section, "sdf_iterations", &opt->sdf_iterations); ReadAttrInt(section, "sdf_initpoints", &opt->sdf_initpoints); diff --git a/src/xml/xml_native_writer.cc b/src/xml/xml_native_writer.cc index 3c610bb3..52ee844c 100644 --- a/src/xml/xml_native_writer.cc +++ b/src/xml/xml_native_writer.cc @@ -940,7 +940,7 @@ void mjXWriter::Option(XMLElement* root) { WriteAttr(section, "tolerance", 1, &model->option.tolerance, &opt.tolerance); WriteAttr(section, "ls_tolerance", 1, &model->option.ls_tolerance, &opt.ls_tolerance); WriteAttr(section, "noslip_tolerance", 1, &model->option.noslip_tolerance, &opt.noslip_tolerance); - WriteAttr(section, "mpr_tolerance", 1, &model->option.mpr_tolerance, &opt.mpr_tolerance); + WriteAttr(section, "ccd_tolerance", 1, &model->option.ccd_tolerance, &opt.ccd_tolerance); WriteAttr(section, "gravity", 3, model->option.gravity, opt.gravity); WriteAttr(section, "wind", 3, model->option.wind, opt.wind); WriteAttr(section, "magnetic", 3, model->option.magnetic, opt.magnetic); @@ -963,7 +963,7 @@ void mjXWriter::Option(XMLElement* root) { WriteAttrInt(section, "iterations", model->option.iterations, opt.iterations); WriteAttrInt(section, "ls_iterations", model->option.ls_iterations, opt.ls_iterations); WriteAttrInt(section, "noslip_iterations", model->option.noslip_iterations, opt.noslip_iterations); - WriteAttrInt(section, "mpr_iterations", model->option.mpr_iterations, opt.mpr_iterations); + WriteAttrInt(section, "ccd_iterations", model->option.ccd_iterations, opt.ccd_iterations); WriteAttrInt(section, "sdf_iterations", model->option.sdf_iterations, opt.sdf_iterations); WriteAttrInt(section, "sdf_initpoints", model->option.sdf_initpoints, opt.sdf_initpoints); diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 08ebfe4e..02747279 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -5008,7 +5008,7 @@ public unsafe struct mjOption_ { public double tolerance; public double ls_tolerance; public double noslip_tolerance; - public double mpr_tolerance; + public double ccd_tolerance; public fixed double gravity[3]; public fixed double wind[3]; public fixed double magnetic[3]; @@ -5025,7 +5025,7 @@ public unsafe struct mjOption_ { public int iterations; public int ls_iterations; public int noslip_iterations; - public int mpr_iterations; + public int ccd_iterations; public int disableflags; public int enableflags; public int disableactuator; diff --git a/unity/Runtime/Components/MjGlobalSettings.cs b/unity/Runtime/Components/MjGlobalSettings.cs index 51a245c5..d41cefe5 100644 --- a/unity/Runtime/Components/MjGlobalSettings.cs +++ b/unity/Runtime/Components/MjGlobalSettings.cs @@ -190,9 +190,9 @@ public struct MjOptionStruct { [Tooltip("Threshold used for early termination of the Noslip solver.")] public float NoSlipTolerance; [Tooltip("Maximum iterations for convex mesh collisions.")] - public int MprIterations; + public int CcdIterations; [Tooltip("Threshold used for early termination of the MPR algorithm.")] - public float MprTolerance; + public float CcdTolerance; public MjcfOptionFlag Flag; @@ -214,8 +214,8 @@ public struct MjOptionStruct { Tolerance = 1e-8f, NoSlipIterations = 0, NoSlipTolerance = 1e-6f, - MprIterations = 50, - MprTolerance = 1e-6f, + CcdIterations = 50, + CcdTolerance = 1e-6f, Flag = MjcfOptionFlag.Default }; @@ -258,8 +258,8 @@ public struct MjOptionStruct { NoSlipIterations = (int)mjcf.GetFloatAttribute( "noslip_iterations", localDefault.NoSlipIterations); NoSlipTolerance = mjcf.GetFloatAttribute("noslip_tolerance", localDefault.NoSlipTolerance); - MprIterations = (int)mjcf.GetFloatAttribute("mpr_iterations", localDefault.MprIterations); - MprTolerance = mjcf.GetFloatAttribute("mpr_tolerance", localDefault.MprTolerance); + CcdIterations = (int)mjcf.GetFloatAttribute("ccd_iterations", localDefault.CcdIterations); + CcdTolerance = mjcf.GetFloatAttribute("ccd_tolerance", localDefault.CcdTolerance); var flagElements = mjcf.GetElementsByTagName("flag"); if (flagElements.Count == 1) { @@ -291,8 +291,8 @@ public struct MjOptionStruct { mjcf.SetAttribute("tolerance", MjEngineTool.MakeLocaleInvariant($"{Tolerance}")); mjcf.SetAttribute("noslip_iterations", MjEngineTool.MakeLocaleInvariant($"{NoSlipIterations}")); mjcf.SetAttribute("noslip_tolerance", MjEngineTool.MakeLocaleInvariant($"{NoSlipTolerance}")); - mjcf.SetAttribute("mpr_iterations", MjEngineTool.MakeLocaleInvariant($"{MprIterations}")); - mjcf.SetAttribute("mpr_tolerance", MjEngineTool.MakeLocaleInvariant($"{MprTolerance}")); + mjcf.SetAttribute("ccd_iterations", MjEngineTool.MakeLocaleInvariant($"{CcdIterations}")); + mjcf.SetAttribute("ccd_tolerance", MjEngineTool.MakeLocaleInvariant($"{CcdTolerance}")); var flags = (XmlElement)mjcf.AppendChild( mjcf.OwnerDocument.CreateElement("flag")); diff --git a/unity/Tests/Editor/Components/MjGlobalSettingsTests.cs b/unity/Tests/Editor/Components/MjGlobalSettingsTests.cs index 4069f45b..1e567b7c 100644 --- a/unity/Tests/Editor/Components/MjGlobalSettingsTests.cs +++ b/unity/Tests/Editor/Components/MjGlobalSettingsTests.cs @@ -56,8 +56,8 @@ public class MjGlobalSettingsGenerationTests { _settings.GlobalOptions.Tolerance = 3.4f; _settings.GlobalOptions.NoSlipIterations = 5; _settings.GlobalOptions.NoSlipTolerance = 6.7f; - _settings.GlobalOptions.MprIterations = 8; - _settings.GlobalOptions.MprTolerance = 0.9f; + _settings.GlobalOptions.CcdIterations = 8; + _settings.GlobalOptions.CcdTolerance = 0.9f; _settings.GlobalSizes.Memory = "1M"; _settings.GlobalsToMjcf(_root); @@ -77,8 +77,8 @@ public class MjGlobalSettingsGenerationTests { Assert.That(_doc.OuterXml, Does.Contain(@"tolerance=""3.4""")); Assert.That(_doc.OuterXml, Does.Contain(@"noslip_iterations=""5""")); Assert.That(_doc.OuterXml, Does.Contain(@"noslip_tolerance=""6.7""")); - Assert.That(_doc.OuterXml, Does.Contain(@"mpr_iterations=""8""")); - Assert.That(_doc.OuterXml, Does.Contain(@"mpr_tolerance=""0.9""")); + Assert.That(_doc.OuterXml, Does.Contain(@"ccd_iterations=""8""")); + Assert.That(_doc.OuterXml, Does.Contain(@"ccd_tolerance=""0.9""")); Assert.That(_doc.OuterXml, Does.Contain(@"memory=""1M""")); } @@ -147,8 +147,8 @@ public class MjGlobalSettingsParsingTests { _option.SetAttribute("tolerance", "3.4"); _option.SetAttribute("noslip_iterations", "5"); _option.SetAttribute("noslip_tolerance", "6.7"); - _option.SetAttribute("mpr_iterations", "8"); - _option.SetAttribute("mpr_tolerance", "0.9"); + _option.SetAttribute("ccd_iterations", "8"); + _option.SetAttribute("ccd_tolerance", "0.9"); _flag.SetAttribute("gravity", "disable"); @@ -172,8 +172,8 @@ public class MjGlobalSettingsParsingTests { Assert.That(_settings.GlobalOptions.Tolerance, Is.EqualTo(3.4f)); Assert.That(_settings.GlobalOptions.NoSlipIterations, Is.EqualTo(5)); Assert.That(_settings.GlobalOptions.NoSlipTolerance, Is.EqualTo(6.7f)); - Assert.That(_settings.GlobalOptions.MprIterations, Is.EqualTo(8)); - Assert.That(_settings.GlobalOptions.MprTolerance, Is.EqualTo(0.9f)); + Assert.That(_settings.GlobalOptions.CcdIterations, Is.EqualTo(8)); + Assert.That(_settings.GlobalOptions.CcdTolerance, Is.EqualTo(0.9f)); Assert.That(_settings.GlobalOptions.Flag.Gravity, Is.EqualTo(EnableDisableFlag.disable)); diff --git a/unity/Tests/Editor/Components/MjcfGenerationContextTests.cs b/unity/Tests/Editor/Components/MjcfGenerationContextTests.cs index dfc908c0..40b534d7 100644 --- a/unity/Tests/Editor/Components/MjcfGenerationContextTests.cs +++ b/unity/Tests/Editor/Components/MjcfGenerationContextTests.cs @@ -74,8 +74,8 @@ public class MjcfGenerationContextTests { Assert.That(mjcf.OuterXml, Does.Contain("tolerance")); Assert.That(mjcf.OuterXml, Does.Contain("noslip_iterations")); Assert.That(mjcf.OuterXml, Does.Contain("noslip_tolerance")); - Assert.That(mjcf.OuterXml, Does.Contain("mpr_iterations")); - Assert.That(mjcf.OuterXml, Does.Contain("mpr_tolerance")); + Assert.That(mjcf.OuterXml, Does.Contain("ccd_iterations")); + Assert.That(mjcf.OuterXml, Does.Contain("ccd_tolerance")); } [Test] From 505d01a1c1018bcbf79f789f5e84f5b8979cb70f Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Fri, 23 Aug 2024 07:17:26 -0700 Subject: [PATCH 23/26] Add error handling to mj_recompile. This CL adds error handling to the mj_recompile function. If the recompile fails, the function will return -1 and set the error message in the spec. The Python wrapper for mj_recompile will catch the error and raise a ValueError exception. PiperOrigin-RevId: 666784579 Change-Id: I225eca1769ea839c782be3c03fc4ff1ea5885a48 --- doc/APIreference/functions.rst | 10 +++++----- doc/APIreference/functions_override.rst | 12 ++++++++---- doc/includes/references.h | 2 +- include/mujoco/mujoco.h | 4 ++-- introspect/functions.py | 4 ++-- python/mujoco/specs_test.py | 24 ++++++++++++++++++++++++ python/mujoco/structs.cc | 4 +++- src/user/user_api.cc | 12 +++++++++--- src/user/user_api.h | 4 ++-- test/user/user_api_test.cc | 21 +++++++++++++++++++++ 10 files changed, 77 insertions(+), 20 deletions(-) diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 96ade764..09619b39 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -47,7 +47,9 @@ mj_compile .. mujoco-include:: mj_compile -Compile spec to model. +Compile :ref:`mjSpec` to :ref:`mjModel`. A spec can be edited and compiled multiple times, returning a new +:ref:`mjModel` instance that takes the edits into account. +If compilation fails, :ref:`mj_compile` returns ``NULL``; the error can be read with :ref:`mjs_getError`. .. _mj_recompile: @@ -63,10 +65,8 @@ reallocate existing :ref:`mjModel` and :ref:`mjData` instances in-place. Second, newly added or removed degrees of freedom. This allows the user to continue simulation with the same model and data struct pointers while editing the model programmatically. -.. admonition:: Incomplete implementation - :class: attention - - This function is currently incomplete, preserving only ``mjData.qpos`` and ``mjData.qvel``. +:ref:`mj_recompile` returns 0 if compilation succeed. In the case of failure, the given :ref:`mjModel` and :ref:`mjData` +instances will be deleted; as in :ref:`mj_compile`, the compilation error can be read with :ref:`mjs_getError`. .. _mj_saveLastXML: diff --git a/doc/APIreference/functions_override.rst b/doc/APIreference/functions_override.rst index 4e7b6b50..d703cf94 100644 --- a/doc/APIreference/functions_override.rst +++ b/doc/APIreference/functions_override.rst @@ -29,6 +29,12 @@ The key function here is :ref:`mj_loadXML`. It invokes the built-in parser and c a valid mjModel, or NULL - in which case the user should check the error information in the user-provided string. The model and all files referenced in it can be loaded from disk or from a VFS when provided. +.. _mj_compile: + +Compile :ref:`mjSpec` to :ref:`mjModel`. A spec can be edited and compiled multiple times, returning a new +:ref:`mjModel` instance that takes the edits into account. +If compilation fails, :ref:`mj_compile` returns ``NULL``; the error can be read with :ref:`mjs_getError`. + .. _mj_recompile: Recompile spec to model, preserving the state. Like :ref:`mj_compile`, this function compiles an :ref:`mjSpec` to an @@ -38,10 +44,8 @@ reallocate existing :ref:`mjModel` and :ref:`mjData` instances in-place. Second, newly added or removed degrees of freedom. This allows the user to continue simulation with the same model and data struct pointers while editing the model programmatically. -.. admonition:: Incomplete implementation - :class: attention - - This function is currently incomplete, preserving only ``mjData.qpos`` and ``mjData.qvel``. +:ref:`mj_recompile` returns 0 if compilation succeed. In the case of failure, the given :ref:`mjModel` and :ref:`mjData` +instances will be deleted; as in :ref:`mj_compile`, the compilation error can be read with :ref:`mjs_getError`. .. _Mainsimulation: diff --git a/doc/includes/references.h b/doc/includes/references.h index a70bc244..4ecd4b74 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -3103,7 +3103,7 @@ mjModel* mj_loadXML(const char* filename, const mjVFS* vfs, char* error, int err mjSpec* mj_parseXML(const char* filename, const mjVFS* vfs, char* error, int error_sz); mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, int error_sz); mjModel* mj_compile(mjSpec* s, const mjVFS* vfs); -void mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d); +int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d); int mj_saveLastXML(const char* filename, const mjModel* m, char* error, int error_sz); void mj_freeLastXML(void); void mj_copyBack(mjSpec* s, const mjModel* m); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 26014657..0d115c9f 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -106,8 +106,8 @@ MJAPI mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, // Compile spec to model. MJAPI mjModel* mj_compile(mjSpec* s, const mjVFS* vfs); -// Recompile spec to model, preserving the state. -MJAPI void mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d); +// Recompile spec to model, preserving the state, return 0 on success. +MJAPI int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d); // Update XML data structures with info from low-level model, save as MJCF. // If error is not NULL, it must have size error_sz. diff --git a/introspect/functions.py b/introspect/functions.py index d787bf1c..90e744a0 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -251,7 +251,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ('mj_recompile', FunctionDecl( name='mj_recompile', - return_type=ValueType(name='void'), + return_type=ValueType(name='int'), parameters=( FunctionParameterDecl( name='s', @@ -278,7 +278,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), ), ), - doc='Recompile spec to model, preserving the state.', + doc='Recompile spec to model, preserving the state, return 0 on success.', # pylint: disable=line-too-long )), ('mj_saveLastXML', FunctionDecl( diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index a860f443..c440d91d 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -352,5 +352,29 @@ class SpecsTest(absltest.TestCase): self.assertEqual(model.nplugin, 1) self.assertEqual(model.body_plugin[1], 0) + def test_recompile_error(self): + main_xml = """ + + + + + + + + """ + + spec = mujoco.MjSpec() + spec.from_string(main_xml) + model = spec.compile() + data = mujoco.MjData(model) + + spec.add_material().name = 'yellow' + spec.add_material().name = 'yellow' + + with self.assertRaisesRegex( + ValueError, "Error: repeated name 'yellow' in material" + ): + spec.recompile(model, data) + if __name__ == '__main__': absltest.main() diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index 66566fe7..20b6678a 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -456,7 +456,9 @@ py::tuple RecompileSpec(raw::MjSpec* spec, const MjModelWrapper& old_m, raw::MjModel* m = static_cast(mju_malloc(sizeof(mjModel))); m->buffer = nullptr; raw::MjData* d = mj_copyData(nullptr, old_m.get(), old_d.get()); - mj_recompile(spec, nullptr, m, d); + if (mj_recompile(spec, nullptr, m, d)) { + throw py::value_error(mjs_getError(spec)); + } py::object m_pyobj = py::cast((MjModelWrapper(m))); py::object d_pyobj = diff --git a/src/user/user_api.cc b/src/user/user_api.cc index f598db61..228299f1 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -90,8 +90,8 @@ mjModel* mj_compile(mjSpec* s, const mjVFS* vfs) { -// recompile spec into existing model and data while preserving the state -void mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d) { +// recompile spec to model, preserving the state, return 0 on success +[[nodiscard]] int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d) { mjCModel* modelC = static_cast(s->element); std::string state_name = "state"; mjtNum time = 0; @@ -99,13 +99,19 @@ void mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d) { time = d->time; modelC->SaveState(state_name, d->qpos, d->qvel, d->act, d->ctrl, d->mocap_pos, d->mocap_quat); } - modelC->Compile(vfs, &m); + if (!modelC->Compile(vfs, &m)) { + if (d) { + mj_deleteData(d); + } + return -1; + }; if (d) { modelC->MakeData(m, &d); modelC->RestoreState(state_name, m->qpos0, m->body_pos, m->body_quat, d->qpos, d->qvel, d->act, d->ctrl, d->mocap_pos, d->mocap_quat); d->time = time; } + return 0; } diff --git a/src/user/user_api.h b/src/user/user_api.h index c30abd53..82ce193a 100644 --- a/src/user/user_api.h +++ b/src/user/user_api.h @@ -42,8 +42,8 @@ MJAPI mjSpec* mj_makeSpec(void); // Compile spec to model. MJAPI mjModel* mj_compile(mjSpec* s, const mjVFS* vfs); -// Recompile spec to model preserving the current state. -MJAPI void mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d); +// Recompile spec to model, preserving the state, return 0 on success. +MJAPI int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d); // Copy spec. MJAPI mjSpec* mj_copySpec(const mjSpec* s); diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 346f179d..cc88aec9 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -140,6 +140,27 @@ TEST_F(PluginTest, ActivatePlugin) { mj_deleteModel(model); } +TEST_F(MujocoTest, RecompileFails) { + mjSpec* spec = mj_makeSpec(); + mjsBody* body = mjs_addBody(mjs_findBody(spec, "world"), 0); + mjsGeom* geom = mjs_addGeom(body, 0); + geom->type = mjGEOM_SPHERE; + geom->size[0] = 1; + + mjModel* model = mj_compile(spec, 0); + mjData* data = mj_makeData(model); + + mjsMaterial* mat1 = mjs_addMaterial(spec, 0); + mjsMaterial* mat2 = mjs_addMaterial(spec, 0); + mjs_setString(mat1->name, "yellow"); + mjs_setString(mat2->name, "yellow"); + + EXPECT_EQ(mj_recompile(spec, 0, model, data), -1); + EXPECT_STREQ(mjs_getError(spec), "Error: repeated name 'yellow' in material"); + + mj_deleteSpec(spec); +} + // ------------------- test recompilation multiple files ----------------------- TEST_F(PluginTest, RecompileCompare) { mjtNum tol = 0; From d360ea1ce6138707c4cb03e570f3f6a4516f4f7e Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Fri, 23 Aug 2024 07:43:43 -0700 Subject: [PATCH 24/26] Minor GJK + EPA improvements: 1. Correct stopping criteria for GJK. 2. Store witness points from EPA in mjCCDObj, 3. Get dir directly from EPA instead of computing dir from witness points. PiperOrigin-RevId: 666791852 Change-Id: I6a168f635b57c17692d31da6f6076b48c32ff2e3 --- src/engine/engine_collision_gjk.c | 40 +++++++++++++++--------- test/engine/engine_collision_gjk_test.cc | 7 +++-- 2 files changed, 29 insertions(+), 18 deletions(-) diff --git a/src/engine/engine_collision_gjk.c b/src/engine/engine_collision_gjk.c index daab5889..ee7dfdcc 100644 --- a/src/engine/engine_collision_gjk.c +++ b/src/engine/engine_collision_gjk.c @@ -97,8 +97,8 @@ static void attachFace(Polytope* pt, int v1, int v2, int v3); // returns the penetration depth (negative distance) of the convex objects // witness points are stored in x1 and x2 -static mjtNum epa(const mjCCDConfig* config, Polytope* pt, - mjCCDObj* obj1, mjCCDObj* obj2, mjtNum x1[3], mjtNum x2[3]); +static mjtNum epa(const mjCCDConfig* config, Polytope* pt, mjCCDObj* obj1, mjCCDObj* obj2, + mjtNum dir[3]); // internal data structure for the returning simplex from GJK typedef struct { @@ -118,6 +118,7 @@ static mjtNum _gjk(const mjCCDConfig* config, mjCCDObj* obj1, mjCCDObj* obj2, mjtNum* x1_k = obj1->x0; mjtNum* x2_k = obj2->x0; mju_sub3(x_k, x1_k, x2_k); + mjtNum epsilon = config->tolerance * config->tolerance; int N = config->max_iterations; for (size_t k = 0; k < N; k++) { @@ -130,10 +131,10 @@ static mjtNum _gjk(const mjCCDConfig* config, mjCCDObj* obj1, mjCCDObj* obj2, mju_sub3(s_k, s1, s2); // the stopping criteria relies on the Frank-Wolfe duality gap given by - // f(x_k) - f(x_min) <= < grad f(x_k), (x_k - s_k) > + // |f(x_k) - f(x_min)|^2 <= < grad f(x_k), (x_k - s_k) > mjtNum diff[3]; mju_sub3(diff, x_k, s_k); - if (2*mju_dot3(x_k, diff) < config->tolerance) { + if (2*mju_dot3(x_k, diff) < epsilon) { break; } @@ -173,6 +174,11 @@ static mjtNum _gjk(const mjCCDConfig* config, mjCCDObj* obj1, mjCCDObj* obj2, // simplex in Minkowski difference mju_copy3(simplex + 3*n++, simplex + 3*i); } + + // we have a tetrahedron containing the origin so return early + if (n == 4) { + break; + } } if (ret1 && ret2) { ret1->nverts = n; @@ -888,8 +894,8 @@ static void epa_witness(const Polytope* pt, int index, mjtNum x1[3], mjtNum x2[3 } // returns the penetration depth (negative distance) of the convex objects -static mjtNum epa(const mjCCDConfig* config, Polytope* pt, - mjCCDObj* obj1, mjCCDObj* obj2, mjtNum x1[3], mjtNum x2[3]) { +static mjtNum epa(const mjCCDConfig* config, Polytope* pt, mjCCDObj* obj1, mjCCDObj* obj2, + mjtNum dir[3]) { mjtNum dist = mjMAXVAL; int index; Horizon h; @@ -941,7 +947,9 @@ static mjtNum epa(const mjCCDConfig* config, Polytope* pt, h.n = 0; // clear horizon } mju_free(h.edges); - epa_witness(pt, index, x1, x2); + Face face = pt->faces[index]; + mju_copy3(dir, face.v); + epa_witness(pt, index, obj1->x0, obj2->x0); return dist; } @@ -949,7 +957,7 @@ static mjtNum epa(const mjCCDConfig* config, Polytope* pt, // runs both GJK and EPA (if needed) static mjtNum _gjk_epa(const mjCCDConfig* config, mjCCDObj* obj1, mjCCDObj* obj2, Polytope* pt, - mjtNum x1[3], mjtNum x2[3]) { + mjtNum dir[3]) { Simplex simplex1, simplex2; mjtNum dist = _gjk(config, obj1, obj2, &simplex1, &simplex2); @@ -965,7 +973,7 @@ static mjtNum _gjk_epa(const mjCCDConfig* config, mjCCDObj* obj1, mjCCDObj* obj2 // simplex not on boundary (objects are penetrating) if (ret) { - dist = epa(config, pt, obj1, obj2, x1, x2); + dist = epa(config, pt, obj1, obj2, dir); return -dist; } return 0; @@ -981,21 +989,23 @@ int mj_gjkPenetration(const void *obj1, const void *obj2, const ccd_t *ccd, Polytope pt; initPolytope(&pt); mjCCDConfig config; + mjCCDObj* o1 = (mjCCDObj*) obj1; + mjtNum* x1 = o1->x0; mjCCDObj* o2 = (mjCCDObj*) obj2; - o1->center(o1->x0, o1); - o2->center(o2->x0, o2); + mjtNum* x2 = o2->x0; + + o1->center(x1, o1); + o2->center(x2, o2); config.max_iterations = ccd->max_iterations; config.tolerance = ccd->mpr_tolerance; - mjtNum x1[3], x2[3]; - mjtNum dist = _gjk_epa(&config, o1, o2, &pt, x1, x2); + mjtNum d[3]; + mjtNum dist = _gjk_epa(&config, o1, o2, &pt, d); if (dist < 0) { if (depth) *depth = -dist; if (dir) { - mjtNum d[3]; - mju_sub3(d, x1, x2); mju_normalize3(d); mju_copy3(dir->v, d); } diff --git a/test/engine/engine_collision_gjk_test.cc b/test/engine/engine_collision_gjk_test.cc index da1d1602..d41972a8 100644 --- a/test/engine/engine_collision_gjk_test.cc +++ b/test/engine/engine_collision_gjk_test.cc @@ -82,10 +82,11 @@ mjtNum run_gjkPenetration(mjModel* m, mjData* d, int g1, int g2, ccd_real_t depth; ccd_vec3_t ccd_dir, ccd_pos; - mj_gjkPenetration(&obj1, &obj2, &ccd, &depth, &ccd_dir, &ccd_pos); + int ret = mj_gjkPenetration(&obj1, &obj2, &ccd, &depth, &ccd_dir, &ccd_pos); + if (ret) return 0; // objects not colliding if (dir) mju_copy3(dir, ccd_dir.v); if (pos) mju_copy3(pos, ccd_pos.v); - return depth; + return -depth; } using MjGjkTest = MujocoTest; @@ -170,7 +171,7 @@ TEST_F(MjGjkTest, BoxBoxIntersect) { mjtNum dir[3], pos[3]; mjtNum dist = run_gjkPenetration(model, data, geom1, geom2, dir, pos); - EXPECT_NEAR(dist, 1, kTolerance); + EXPECT_NEAR(dist, -1, kTolerance); EXPECT_NEAR(dir[0], 1, kTolerance); EXPECT_NEAR(dir[1], 0, kTolerance); EXPECT_NEAR(dir[2], 0, kTolerance); From 9d4f8edad9d2b11a31aaf0483547a6d9530cbb89 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Fri, 23 Aug 2024 08:12:52 -0700 Subject: [PATCH 25/26] Avoid pointer access in loop conditionals in engine_core_smooth.c PiperOrigin-RevId: 666800203 Change-Id: Iec6ce752d0cac1a98809ab085ed8f94416f79e95 --- src/engine/engine_core_smooth.c | 68 +++++++++++++++++++-------------- 1 file changed, 39 insertions(+), 29 deletions(-) diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 57afead3..2fec9331 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -35,6 +35,8 @@ // forward kinematics void mj_kinematics(const mjModel* m, mjData* d) { + int nbody = m->nbody, nsite = m->nsite, ngeom = m->ngeom; + // set world position and orientation mju_zero3(d->xpos); mju_unit4(d->xquat); @@ -45,7 +47,7 @@ void mj_kinematics(const mjModel* m, mjData* d) { d->ximat[0] = d->ximat[4] = d->ximat[8] = 1; // compute global cartesian positions and orientations of all bodies - for (int i=1; i < m->nbody; i++) { + for (int i=1; i < nbody; i++) { mjtNum xpos[3], xquat[4]; int jntadr = m->body_jntadr[i]; int jntnum = m->body_jntnum[i]; @@ -153,21 +155,21 @@ void mj_kinematics(const mjModel* m, mjData* d) { } // compute/copy Cartesian positions and orientations of body inertial frames - for (int i=1; i < m->nbody; i++) { + for (int i=1; i < nbody; i++) { mj_local2Global(d, d->xipos+3*i, d->ximat+9*i, m->body_ipos+3*i, m->body_iquat+4*i, i, m->body_sameframe[i]); } // compute/copy Cartesian positions and orientations of geoms - for (int i=0; i < m->ngeom; i++) { + for (int i=0; i < ngeom; i++) { mj_local2Global(d, d->geom_xpos+3*i, d->geom_xmat+9*i, m->geom_pos+3*i, m->geom_quat+4*i, m->geom_bodyid[i], m->geom_sameframe[i]); } // compute/copy Cartesian positions and orientations of sites - for (int i=0; i < m->nsite; i++) { + for (int i=0; i < nsite; i++) { mj_local2Global(d, d->site_xpos+3*i, d->site_xmat+9*i, m->site_pos+3*i, m->site_quat+4*i, m->site_bodyid[i], m->site_sameframe[i]); @@ -178,6 +180,7 @@ void mj_kinematics(const mjModel* m, mjData* d) { // map inertias and motion dofs to global frame centered at subtree-CoM void mj_comPos(const mjModel* m, mjData* d) { + int nbody = m->nbody, njnt = m->njnt; mjtNum offset[3], axis[3]; mj_markStack(d); mjtNum* mass_subtree = mj_stackAllocNum(d, m->nbody); @@ -187,7 +190,7 @@ void mj_comPos(const mjModel* m, mjData* d) { mju_zero(d->subtree_com, m->nbody*3); // backwards pass over bodies: compute subtree_com and mass_subtree - for (int i=m->nbody-1; i >= 0; i--) { + for (int i=nbody-1; i >= 0; i--) { // add local info mju_addToScl3(d->subtree_com+3*i, d->xipos+3*i, m->body_mass[i]); mass_subtree[i] += m->body_mass[i]; @@ -212,14 +215,14 @@ void mj_comPos(const mjModel* m, mjData* d) { mju_zero(d->cinert, 10); // map inertias to frame centered at subtree_com - for (int i=1; i < m->nbody; i++) { + for (int i=1; i < nbody; i++) { mju_sub3(offset, d->xipos+3*i, d->subtree_com+3*m->body_rootid[i]); mju_inertCom(d->cinert+10*i, m->body_inertia+3*i, d->ximat+9*i, offset, m->body_mass[i]); } // map motion dofs to global frame centered at subtree_com - for (int j=0; j < m->njnt; j++) { + for (int j=0; j < njnt; j++) { // get dof address, body index int da = 6*m->jnt_dofadr[j]; int bi = m->jnt_bodyid[j]; @@ -471,7 +474,8 @@ void mj_flex(const mjModel* m, mjData* d) { int dim = m->flex_dim[f]; // process elements of this flex - for (int e=0; e < m->flex_elemnum[f]; e++) { + int elemnum = m->flex_elemnum[f]; + for (int e=0; e < elemnum; e++) { const int* edata = m->flex_elem + m->flex_elemdataadr[f] + e*(dim+1); const mjtNum* vert = d->flexvert_xpos + 3*m->flex_vertadr[f]; @@ -652,6 +656,7 @@ void mj_tendon(const mjModel* m, mjData* d) { adr = m->tendon_adr[i]; d->ten_wrapadr[i] = wcnt; d->ten_wrapnum[i] = 0; + int tendon_num = m->tendon_num[i]; // sparse Jacobian row init if (issparse) { @@ -661,7 +666,7 @@ void mj_tendon(const mjModel* m, mjData* d) { // process joint tendon if (m->wrap_type[adr] == mjWRAP_JOINT) { // process all defined joints - for (int j=0; j < m->tendon_num[i]; j++) { + for (int j=0; j < tendon_num; j++) { // get joint id int k = m->wrap_objid[adr+j]; @@ -683,10 +688,10 @@ void mj_tendon(const mjModel* m, mjData* d) { // sort on colind if sparse: custom insertion sort if (issparse) { - int x, *list = colind+rowadr[i]; + int x, *list = colind+rowadr[i], nnz = rownnz[i]; mjtNum y, *listy = J+rowadr[i]; - for (int k=1; k < rownnz[i]; k++) { + for (int k=1; k < nnz; k++) { x = list[k]; y = listy[k]; int j = k-1; @@ -706,7 +711,7 @@ void mj_tendon(const mjModel* m, mjData* d) { // process spatial tendon divisor = 1; int j = 0; - while (j < m->tendon_num[i]-1) { + while (j < tendon_num-1) { // get 1st and 2nd object tp0 = m->wrap_type[adr+j]; id0 = m->wrap_objid[adr+j]; @@ -824,7 +829,7 @@ void mj_tendon(const mjModel* m, mjData* d) { j += (tpw != mjWRAP_NONE ? 2 : 1); // assign last site before pulley or tendon end - if (j == m->tendon_num[i]-1 || m->wrap_type[adr+j+1] == mjWRAP_PULLEY) { + if (j == tendon_num-1 || m->wrap_type[adr+j+1] == mjWRAP_PULLEY) { mju_copy3(d->wrap_xpos+wcnt*3, d->site_xpos+3*id1); d->wrap_obj[wcnt] = -1; d->ten_wrapnum[i]++; @@ -1164,8 +1169,8 @@ void mj_transmission(const mjModel* m, mjData* d) { mju_zero(moment_exclude, nv); // count all relevant contacts, accumulate Jacobians - int counter = 0; - for (int j=0; j < d->ncon; j++) { + int counter = 0, ncon = d->ncon; + for (int j=0; j < ncon; j++) { const mjContact* con = d->contact+j; // get geom ids @@ -1618,13 +1623,14 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) // compute cvel, cdof_dot void mj_comVel(const mjModel* m, mjData* d) { + int nbody = m->nbody; mjtNum tmp[6], cvel[6], cdofdot[36]; // set world vel to 0 mju_zero(d->cvel, 6); // forward pass over bodies - for (int i=1; i < m->nbody; i++) { + for (int i=1; i < nbody; i++) { // get body's first dof address int bda = m->body_dofadr[i]; @@ -1632,7 +1638,8 @@ void mj_comVel(const mjModel* m, mjData* d) { mju_copy(cvel, d->cvel+6*m->body_parentid[i], 6); // cvel = cvel_parent + cdof * qvel, cdofdot = cvel x cdof - for (int j=0; j < m->body_dofnum[i]; j++) { + int dofnum = m->body_dofnum[i]; + for (int j=0; j < dofnum; j++) { // compute cvel and cdofdot switch ((mjtJoint) m->jnt_type[m->dof_jntid[bda+j]]) { case mjJNT_FREE: @@ -1683,12 +1690,13 @@ void mj_comVel(const mjModel* m, mjData* d) { // subtree linear velocity and angular momentum void mj_subtreeVel(const mjModel* m, mjData* d) { + int nbody = m->nbody; mjtNum dx[3], dv[3], dp[3], dL[3]; mj_markStack(d); mjtNum* body_vel = mj_stackAllocNum(d, 6*m->nbody); // bodywise quantities - for (int i=0; i < m->nbody; i++) { + for (int i=0; i < nbody; i++) { // compute and save body velocity mj_objectVelocity(m, d, mjOBJ_BODY, i, body_vel+6*i, 0); @@ -1704,7 +1712,7 @@ void mj_subtreeVel(const mjModel* m, mjData* d) { } // subtree linvel - for (int i=m->nbody-1; i >= 0; i--) { + for (int i=nbody-1; i >= 0; i--) { // non-world: add linear momentum to parent if (i) { mju_addTo3(d->subtree_linvel+3*m->body_parentid[i], d->subtree_linvel+3*i); @@ -1716,7 +1724,7 @@ void mj_subtreeVel(const mjModel* m, mjData* d) { } // subtree angmom - for (int i=m->nbody-1; i > 0; i--) { + for (int i=nbody-1; i > 0; i--) { int parent = m->body_parentid[i]; // momentum wrt body i @@ -1749,6 +1757,7 @@ void mj_subtreeVel(const mjModel* m, mjData* d) { // RNE: compute M(qpos)*qacc + C(qpos,qvel); flg_acc=0 removes inertial term void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result) { + int nbody = m->nbody, nv = m->nv; mjtNum tmp[6], tmp1[6]; mj_markStack(d); mjtNum* loc_cacc = mj_stackAllocNum(d, m->nbody*6); @@ -1761,7 +1770,7 @@ void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result) { } // forward pass over bodies: accumulate cacc, set cfrc_body - for (int i=1; i < m->nbody; i++) { + for (int i=1; i < nbody; i++) { // get body's first dof address int bda = m->body_dofadr[i]; @@ -1786,13 +1795,13 @@ void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result) { mju_zero(loc_cfrc_body, 6); // backward pass over bodies: accumulate cfrc_body from children - for (int i=m->nbody-1; i > 0; i--) + for (int i=nbody-1; i > 0; i--) if (m->body_parentid[i]) { mju_addTo(loc_cfrc_body+6*m->body_parentid[i], loc_cfrc_body+6*i, 6); } // result = cdof * cfrc_body - for (int i=0; i < m->nv; i++) { + for (int i=0; i < nv; i++) { result[i] = mju_dot(d->cdof+6*i, loc_cfrc_body+6*m->dof_bodyid[i], 6); } @@ -1803,7 +1812,7 @@ void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result) { // RNE with complete data: compute cacc, cfrc_ext, cfrc_int void mj_rnePostConstraint(const mjModel* m, mjData* d) { - int nbody=m->nbody; + int nbody = m->nbody; mjtNum cfrc_com[6], cfrc[6], lfrc[6]; mjContact* con; @@ -1829,7 +1838,8 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { } // cfrc_ext += contacts - for (int i=0; i < d->ncon; i++) + int ncon = d->ncon; + for (int i=0; i < ncon; i++) if (d->contact[i].efc_address >= 0) { // get contact pointer con = d->contact+i; @@ -1867,8 +1877,8 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { } // cfrc_ext += connect and weld constraints - int i = 0; - while (i < d->ne) { + int i = 0, ne = d->ne; + while (i < ne) { if (d->efc_type[i] != mjCNSTR_EQUALITY) mjERROR("row %d of efc is not an equality constraint", i); // SHOULD NOT OCCUR @@ -1942,7 +1952,7 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { // forward pass over bodies: compute cacc, cfrc_int mjtNum cacc[6], cfrc_body[6], cfrc_corr[6]; mju_zero(d->cfrc_int, 6); - for (int j=1; j < m->nbody; j++) { + for (int j=1; j < nbody; j++) { // get body's first dof address int bda = m->body_dofadr[j]; @@ -1963,7 +1973,7 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) { } // backward pass over bodies: accumulate cfrc_int from children - for (int j=m->nbody-1; j > 0; j--) { + for (int j=nbody-1; j > 0; j--) { mju_addTo(d->cfrc_int+6*m->body_parentid[j], d->cfrc_int+6*j, 6); } } From cad8ca6ffd269668c02b526b602b729adf91cdc8 Mon Sep 17 00:00:00 2001 From: Kyle Bayes Date: Fri, 23 Aug 2024 08:36:46 -0700 Subject: [PATCH 26/26] Do pre-checking when creating initial polytope for EPA. PiperOrigin-RevId: 666807572 Change-Id: I64c3715dc72c24f9b6b2206b3568951bd8c2384a --- src/engine/engine_collision_gjk.c | 175 +++++++++++++++++++----------- 1 file changed, 111 insertions(+), 64 deletions(-) diff --git a/src/engine/engine_collision_gjk.c b/src/engine/engine_collision_gjk.c index ee7dfdcc..9138e31c 100644 --- a/src/engine/engine_collision_gjk.c +++ b/src/engine/engine_collision_gjk.c @@ -19,6 +19,7 @@ #include #include +#include #include "engine/engine_collision_convex.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" @@ -618,16 +619,12 @@ static void rotmat(mjtNum R[9], const mjtNum axis[3]) { // creates a polytope from a 1-simplex (2 points i.e. line segment) static int polytope2(Polytope* pt, const mjtNum simplex1[6], const mjtNum simplex2[6], mjCCDObj* obj1, mjCCDObj* obj2) { - const mjtNum* s1a = simplex1; - const mjtNum* s1b = simplex2; - const mjtNum* s2a = simplex1 + 3; - const mjtNum* s2b = simplex2 + 3; - mjtNum s1[3], s2[3]; - mju_sub3(s1, s1a, s1b); - mju_sub3(s2, s2a, s2b); + mjtNum v1[3], v2[3]; + mju_sub3(v1, simplex1 + 0, simplex2 + 0); + mju_sub3(v2, simplex1 + 3, simplex2 + 3); mjtNum diff[3]; - mju_sub3(diff, s2, s1); + mju_sub3(diff, v2, v1); // find component with smallest magnitude (so cross product is largest) mjtNum value = mjMAXVAL; @@ -653,31 +650,50 @@ static int polytope2(Polytope* pt, const mjtNum simplex1[6], const mjtNum simple mju_mulMatVec(d3, R, d2, 3, 3); - mjtNum v1a[3], v2a[3], v3a[3]; - mjtNum v1b[3], v2b[3], v3b[3]; - mjtNum v1[3], v2[3], v3[3]; - support(v1a, v1b, obj1, obj2, d1); - support(v2a, v2b, obj1, obj2, d2); - support(v3a, v3b, obj1, obj2, d3); - - mju_sub3(v1, v1a, v1b); - mju_sub3(v2, v2a, v2b); + mjtNum v3a[3], v3b[3], v3[3]; + support(v3a, v3b, obj1, obj2, d1); mju_sub3(v3, v3a, v3b); + mjtNum v4a[3], v4b[3], v4[3]; + support(v4a, v4b, obj1, obj2, d2); + mju_sub3(v4, v4a, v4b); - int s1i = newVertex(pt, s1a, s1b); - int v1i = newVertex(pt, v1a, v1b); - int v2i = newVertex(pt, v2a, v2b); + mjtNum v5a[3], v5b[3], v5[3]; + support(v5a, v5b, obj1, obj2, d3); + mju_sub3(v5, v5a, v5b); + + // check that all six faces are valid triangles (not collinear) + if (mju_abs(det3(v1, v3, v4)) < mjMINVAL || mju_abs(det3(v1, v3, v5)) < mjMINVAL || + mju_abs(det3(v1, v3, v5)) < mjMINVAL || mju_abs(det3(v2, v3, v4)) < mjMINVAL || + mju_abs(det3(v2, v3, v5)) < mjMINVAL || mju_abs(det3(v2, v4, v5)) < mjMINVAL) { + return 0; + } + + // save vertices and get indices for each one + int v1i = newVertex(pt, simplex1 + 0, simplex2 + 0); + int v2i = newVertex(pt, simplex1 + 3, simplex2 + 3); int v3i = newVertex(pt, v3a, v3b); - int s2i = newVertex(pt, s2a, s2b); + int v4i = newVertex(pt, v4a, v4b); + int v5i = newVertex(pt, v5a, v5b); - // TODO(kylebayes): check what side of the hexahedron the origin is on - attachFace(pt, s1i, v2i, v1i); - attachFace(pt, s1i, v3i, v1i); - attachFace(pt, s1i, v3i, v2i); - attachFace(pt, s2i, v1i, v2i); - attachFace(pt, s2i, v1i, v3i); - attachFace(pt, s2i, v2i, v3i); + + // build hexahedron + attachFace(pt, v1i, v3i, v4i); + attachFace(pt, v1i, v3i, v5i); + attachFace(pt, v1i, v4i, v5i); + attachFace(pt, v2i, v3i, v4i); + attachFace(pt, v2i, v3i, v5i); + attachFace(pt, v2i, v4i, v5i); + + // if the origin is on the affine hull of any of the faces then the origin is not in the + // hexahedron or the hexahedron is degenerate + for (int i = 0; i < 6; i++) { + if (pt->faces[i].dist < mjMINVAL) { + return 0; + } + } + + // valid hexahedron for EPA return 1; } @@ -686,46 +702,70 @@ static int polytope2(Polytope* pt, const mjtNum simplex1[6], const mjtNum simple // creates a polytope from a 2-simplex (3 points i.e. triangle) static int polytope3(Polytope* pt, const mjtNum simplex1[9], const mjtNum simplex2[9], mjCCDObj* obj1, mjCCDObj* obj2) { - const mjtNum* s1a = simplex1; - const mjtNum* s2a = simplex1 + 3; - const mjtNum* s3a = simplex1 + 6; + // get vertices of simplex from GJK + mjtNum v1[3], v2[3], v3[3]; + mju_sub3(v1, simplex1 + 0, simplex2 + 0); + mju_sub3(v2, simplex1 + 3, simplex2 + 3); + mju_sub3(v3, simplex1 + 6, simplex2 + 6); - const mjtNum* s1b = simplex2; - const mjtNum* s2b = simplex2 + 3; - const mjtNum* s3b = simplex2 + 6; - - mjtNum s1[3], s2[3], s3[3]; - mju_sub3(s1, s1a, s1b); - mju_sub3(s2, s2a, s2b); - mju_sub3(s3, s3a, s3b); - - // form hexahedron from triangle and two face normals - - mjtNum diff1[3], diff2[3], n[3], neg_n[3]; - mju_sub3(diff1, s2, s1); - mju_sub3(diff2, s3, s1); + // get normals in both directions + mjtNum diff1[3], diff2[3], n[3], nn[3]; + mju_sub3(diff1, v2, v1); + mju_sub3(diff2, v3, v1); mju_cross(n, diff1, diff2); - mju_scl3(neg_n, n, -1); + if (mju_norm3(n) < mjMINVAL) { + return 0; + } - mjtNum na[3], nb[3], nna[3], nnb[3]; - support(na, nb, obj1, obj2, n); - support(nna, nnb, obj1, obj2, neg_n); + // negative of triangle normal n + mju_scl3(nn, n, -1); - int ni = newVertex(pt, na, nb); - int s1i = newVertex(pt, s1a, s1b); - int s2i = newVertex(pt, s2a, s2b); - int s3i = newVertex(pt, s3a, s3b); - int nni = newVertex(pt, nna, nnb); + // get 4th vertex in n direction + mjtNum v4a[3], v4b[3], v4[3]; + support(v4a, v4b, obj1, obj2, n); + mju_sub3(v4, v4a, v4b); - attachFace(pt, s1i, s2i, ni); - attachFace(pt, s3i, s1i, ni); - attachFace(pt, s2i, s3i, ni); + // we must check that all three faces are valid triangles (not collinear) + if (mju_abs(det3(v4, v1, v2)) < mjMINVAL || + mju_abs(det3(v4, v2, v3)) < mjMINVAL || + mju_abs(det3(v4, v3, v1)) < mjMINVAL) { + return 0; + } - attachFace(pt, s1i, s2i, nni); - attachFace(pt, s3i, s1i, nni); - attachFace(pt, s2i, s3i, nni); + // get 5th vertex in -n direction + mjtNum v5a[3], v5b[3], v5[3]; + support(v5a, v5b, obj1, obj2, nn); + mju_sub3(v5, v5a, v4b); - // TODO(kylebayes): check what side of the hexahedron the origin is on + // we must check that all three faces are valid triangles (not collinear) + if (mju_abs(det3(v5, v1, v2)) < mjMINVAL || + mju_abs(det3(v5, v2, v3)) < mjMINVAL || + mju_abs(det3(v5, v3, v1)) < mjMINVAL) { + return 0; + } + + // save vertices and get indices for each one + int v1i = newVertex(pt, simplex1 + 0, simplex2 + 0); + int v2i = newVertex(pt, simplex1 + 3, simplex2 + 3); + int v3i = newVertex(pt, simplex1 + 6, simplex2 + 6); + int v5i = newVertex(pt, v5a, v5b); + int v4i = newVertex(pt, v4a, v4b); + + // create hexahedron for EPA + attachFace(pt, v1i, v2i, v4i); + attachFace(pt, v3i, v1i, v4i); + attachFace(pt, v2i, v3i, v4i); + attachFace(pt, v1i, v2i, v5i); + attachFace(pt, v3i, v1i, v5i); + attachFace(pt, v2i, v3i, v5i); + + // if the origin is on the affine hull of any of the faces then the origin is not in the + // hexahedron or the hexahedron is degenerate + for (int i = 0; i < 6; i++) { + if (pt->faces[i].dist < mjMINVAL) { + return 0; + } + } return 1; } @@ -733,7 +773,7 @@ static int polytope3(Polytope* pt, const mjtNum simplex1[9], const mjtNum simple // creates a polytope from a 3-simplex (4 points i.e. tetrahedron) static int polytope4(Polytope* pt, const mjtNum simplex1[12], const mjtNum simplex2[12]) { - int v1 = newVertex(pt, simplex1, simplex2); + int v1 = newVertex(pt, simplex1 + 0, simplex2 + 0); int v2 = newVertex(pt, simplex1 + 3, simplex2 + 3); int v3 = newVertex(pt, simplex1 + 6, simplex2 + 6); int v4 = newVertex(pt, simplex1 + 9, simplex2 + 9); @@ -742,8 +782,6 @@ static int polytope4(Polytope* pt, const mjtNum simplex1[12], const mjtNum simpl attachFace(pt, v1, v2, v4); attachFace(pt, v1, v4, v3); attachFace(pt, v4, v2, v3); - - // TODO(kylebayes): check if contains origin return 1; } @@ -893,6 +931,8 @@ static void epa_witness(const Polytope* pt, int index, mjtNum x1[3], mjtNum x2[3 lincomb(x2, lambda, simplex2, 3); } + + // returns the penetration depth (negative distance) of the convex objects static mjtNum epa(const mjCCDConfig* config, Polytope* pt, mjCCDObj* obj1, mjCCDObj* obj2, mjtNum dir[3]) { @@ -906,6 +946,7 @@ static mjtNum epa(const mjCCDConfig* config, Polytope* pt, mjCCDObj* obj1, mjCCD for (int j = 0; j < N; j++) { // find the closest face to the origin dist = mjMAXVAL; + index = -1; for (int i = 0; i < pt->nfaces; i++) { if (pt->faces[i].ignored) continue; if (pt->faces[i].dist < dist) { @@ -914,6 +955,12 @@ static mjtNum epa(const mjCCDConfig* config, Polytope* pt, mjCCDObj* obj1, mjCCD } } + // check if index is set + if (index < 0) { + mju_warning("EPA: empty polytope (most likely a bug)"); + return 0; // assume 0 depth + } + // compute support point w from the closest face's normal mjtNum w1[3], w2[3], w[3]; support(w1, w2, obj1, obj2, pt->faces[index].v);