From 25485ae881f55166b0f35ab395507b4f1b71b3a5 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Fri, 5 Dec 2025 17:40:44 +0100 Subject: [PATCH 001/107] no more xmat Signed-off-by: Alain Denzler --- mujoco_warp/_src/constraint.py | 14 ++++++-------- mujoco_warp/_src/io.py | 17 +++++++++++++++-- mujoco_warp/_src/sensor.py | 32 +++++++++++++++----------------- mujoco_warp/_src/smooth.py | 26 ++++++++++---------------- mujoco_warp/_src/smooth_test.py | 2 -- mujoco_warp/_src/types.py | 2 -- 6 files changed, 46 insertions(+), 47 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 5d270cb8e..81d41a7eb 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -147,7 +147,7 @@ def _efc_equality_connect( qvel_in: wp.array2d(dtype=float), eq_active_in: wp.array2d(dtype=bool), xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), + xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3), cdof_in: wp.array2d(dtype=wp.spatial_vector), @@ -196,8 +196,8 @@ def _efc_equality_connect( else: body1id = obj1id body2id = obj2id - pos1 = xpos_in[worldid, body1id] + xmat_in[worldid, body1id] @ anchor1 - pos2 = xpos_in[worldid, body2id] + xmat_in[worldid, body2id] @ anchor2 + pos1 = xpos_in[worldid, body1id] + math.rot_vec_quat(anchor1, xquat_in[worldid, body1id]) + pos2 = xpos_in[worldid, body2id] + math.rot_vec_quat(anchor2, xquat_in[worldid, body2id]) # error is difference in global positions pos = pos1 - pos2 @@ -740,7 +740,6 @@ def _efc_equality_weld( eq_active_in: wp.array2d(dtype=bool), xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), - xmat_in: wp.array2d(dtype=wp.mat33), site_xpos_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3), cdof_in: wp.array2d(dtype=wp.spatial_vector), @@ -797,8 +796,8 @@ def _efc_equality_weld( else: body1id = obj1id body2id = obj2id - pos1 = xpos_in[worldid, body1id] + xmat_in[worldid, body1id] @ anchor2 - pos2 = xpos_in[worldid, body2id] + xmat_in[worldid, body2id] @ anchor1 + pos1 = xpos_in[worldid, body1id] + math.rot_vec_quat(anchor2, xquat_in[worldid, body1id]) + pos2 = xpos_in[worldid, body2id] + math.rot_vec_quat(anchor1, xquat_in[worldid, body2id]) quat = math.mul_quat(xquat_in[worldid, body1id], relpose) quat1 = math.quat_inv(xquat_in[worldid, body2id]) @@ -1603,7 +1602,7 @@ def make_constraint(m: types.Model, d: types.Data): d.qvel, d.eq_active, d.xpos, - d.xmat, + d.xquat, d.site_xpos, d.subtree_com, d.cdof, @@ -1648,7 +1647,6 @@ def make_constraint(m: types.Model, d: types.Data): d.eq_active, d.xpos, d.xquat, - d.xmat, d.site_xpos, d.subtree_com, d.cdof, diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 098cfab44..fcd8803b3 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -915,8 +915,21 @@ def get_data_into( result.qacc[:] = d.qacc.numpy()[world_id] result.act_dot[:] = d.act_dot.numpy()[world_id] result.xpos[:] = d.xpos.numpy()[world_id] - result.xquat[:] = d.xquat.numpy()[world_id] - result.xmat[:] = d.xmat.numpy()[world_id].reshape((-1, 9)) + xquat = d.xquat.numpy()[world_id] + result.xquat[:] = xquat + # compute xmat from xquat + q = xquat + xmat = np.empty((q.shape[0], 9), dtype=np.float32) + xmat[:, 0] = q[:, 0] ** 2 + q[:, 1] ** 2 - q[:, 2] ** 2 - q[:, 3] ** 2 + xmat[:, 1] = 2.0 * (q[:, 1] * q[:, 2] - q[:, 0] * q[:, 3]) + xmat[:, 2] = 2.0 * (q[:, 1] * q[:, 3] + q[:, 0] * q[:, 2]) + xmat[:, 3] = 2.0 * (q[:, 1] * q[:, 2] + q[:, 0] * q[:, 3]) + xmat[:, 4] = q[:, 0] ** 2 - q[:, 1] ** 2 + q[:, 2] ** 2 - q[:, 3] ** 2 + xmat[:, 5] = 2.0 * (q[:, 2] * q[:, 3] - q[:, 0] * q[:, 1]) + xmat[:, 6] = 2.0 * (q[:, 1] * q[:, 3] - q[:, 0] * q[:, 2]) + xmat[:, 7] = 2.0 * (q[:, 2] * q[:, 3] + q[:, 0] * q[:, 1]) + xmat[:, 8] = q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2 + q[:, 3] ** 2 + result.xmat[:] = xmat result.xipos[:] = d.xipos.numpy()[world_id] result.ximat[:] = d.ximat.numpy()[world_id].reshape((-1, 9)) result.xanchor[:] = d.xanchor.numpy()[world_id] diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 31ba21ffa..44b08f529 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -279,7 +279,7 @@ def _limit_pos( def _frame_pos( # Data in: xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), + xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), ximat_in: wp.array2d(dtype=wp.mat33), geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -316,7 +316,7 @@ def _frame_pos( xmat_ref = ximat_in[worldid, refid] elif objtype == ObjType.XBODY: xpos_ref = xpos_in[worldid, refid] - xmat_ref = xmat_in[worldid, refid] + xmat_ref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: xpos_ref = geom_xpos_in[worldid, refid] xmat_ref = geom_xmat_in[worldid, refid] @@ -337,7 +337,7 @@ def _frame_pos( @wp.func def _frame_axis( # Data in: - xmat_in: wp.array2d(dtype=wp.mat33), + xquat_in: wp.array2d(dtype=wp.quat), ximat_in: wp.array2d(dtype=wp.mat33), geom_xmat_in: wp.array2d(dtype=wp.mat33), site_xmat_in: wp.array2d(dtype=wp.mat33), @@ -354,7 +354,7 @@ def _frame_axis( xmat = ximat_in[worldid, objid] axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) elif objtype == ObjType.XBODY: - xmat = xmat_in[worldid, objid] + xmat = math.quat_to_mat(xquat_in[worldid, objid]) axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) elif objtype == ObjType.GEOM: xmat = geom_xmat_in[worldid, objid] @@ -374,7 +374,7 @@ def _frame_axis( if reftype == ObjType.BODY: xmat_ref = ximat_in[worldid, refid] elif reftype == ObjType.XBODY: - xmat_ref = xmat_in[worldid, refid] + xmat_ref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: xmat_ref = geom_xmat_in[worldid, refid] elif reftype == ObjType.SITE: @@ -491,7 +491,6 @@ def _sensor_pos( qpos_in: wp.array2d(dtype=float), xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), - xmat_in: wp.array2d(dtype=wp.mat33), xipos_in: wp.array2d(dtype=wp.vec3), ximat_in: wp.array2d(dtype=wp.mat33), geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -545,7 +544,7 @@ def _sensor_pos( reftype = sensor_reftype[sensorid] vec3 = _frame_pos( xpos_in, - xmat_in, + xquat_in, xipos_in, ximat_in, geom_xpos_in, @@ -572,7 +571,7 @@ def _sensor_pos( elif sensortype == SensorType.FRAMEZAXIS: axis = 2 vec3 = _frame_axis( - ximat_in, xmat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis + xquat_in, ximat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEQUAT: @@ -855,7 +854,6 @@ def sensor_pos(m: Model, d: Data): d.qpos, d.xpos, d.xquat, - d.xmat, d.xipos, d.ximat, d.geom_xpos, @@ -1051,7 +1049,7 @@ def _frame_linvel( cam_bodyid: wp.array(dtype=int), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), + xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), ximat_in: wp.array2d(dtype=wp.mat33), geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1087,7 +1085,7 @@ def _frame_linvel( xmatref = ximat_in[worldid, refid] elif reftype == ObjType.XBODY: xposref = xpos_in[worldid, refid] - xmatref = xmat_in[worldid, refid] + xmatref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: xposref = geom_xpos_in[worldid, refid] xmatref = geom_xmat_in[worldid, refid] @@ -1157,7 +1155,7 @@ def _frame_angvel( cam_bodyid: wp.array(dtype=int), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), + xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), ximat_in: wp.array2d(dtype=wp.mat33), geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1197,7 +1195,7 @@ def _frame_angvel( if reftype == ObjType.BODY: xmatref = ximat_in[worldid, refid] elif reftype == ObjType.XBODY: - xmatref = xmat_in[worldid, refid] + xmatref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: xmatref = geom_xmat_in[worldid, refid] elif reftype == ObjType.SITE: @@ -1260,7 +1258,7 @@ def _sensor_vel( # Data in: qvel_in: wp.array2d(dtype=float), xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), + xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), ximat_in: wp.array2d(dtype=wp.mat33), geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1312,7 +1310,7 @@ def _sensor_vel( site_bodyid, cam_bodyid, xpos_in, - xmat_in, + xquat_in, xipos_in, ximat_in, geom_xpos_in, @@ -1340,7 +1338,7 @@ def _sensor_vel( site_bodyid, cam_bodyid, xpos_in, - xmat_in, + xquat_in, xipos_in, ximat_in, geom_xpos_in, @@ -1395,7 +1393,7 @@ def sensor_vel(m: Model, d: Data): m.sensor_vel_adr, d.qvel, d.xpos, - d.xmat, + d.xquat, d.xipos, d.ximat, d.geom_xpos, diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index c90784f0d..f570e5529 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -45,7 +45,6 @@ def _kinematics_root( # Data out: xpos_out: wp.array2d(dtype=wp.vec3), xquat_out: wp.array2d(dtype=wp.quat), - xmat_out: wp.array2d(dtype=wp.mat33), xipos_out: wp.array2d(dtype=wp.vec3), ximat_out: wp.array2d(dtype=wp.mat33), ): @@ -53,7 +52,6 @@ def _kinematics_root( xpos_out[worldid, 0] = wp.vec3(0.0) xquat_out[worldid, 0] = wp.quat(1.0, 0.0, 0.0, 0.0) xipos_out[worldid, 0] = wp.vec3(0.0) - xmat_out[worldid, 0] = wp.identity(n=3, dtype=wp.float32) ximat_out[worldid, 0] = wp.identity(n=3, dtype=wp.float32) @@ -79,13 +77,11 @@ def _kinematics_level( mocap_quat_in: wp.array2d(dtype=wp.quat), xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), - xmat_in: wp.array2d(dtype=wp.mat33), # In: body_tree_: wp.array(dtype=int), # Data out: xpos_out: wp.array2d(dtype=wp.vec3), xquat_out: wp.array2d(dtype=wp.quat), - xmat_out: wp.array2d(dtype=wp.mat33), xipos_out: wp.array2d(dtype=wp.vec3), ximat_out: wp.array2d(dtype=wp.mat33), xanchor_out: wp.array2d(dtype=wp.vec3), @@ -129,7 +125,7 @@ def _kinematics_level( body_pos_ = body_pos[body_pos_id, bodyid] body_quat_ = body_quat[body_quat_id, bodyid] - xpos = (xmat_in[worldid, pid] * body_pos_) + xpos_in[worldid, pid] + xpos = math.rot_vec_quat(body_pos_, xquat_in[worldid, pid]) + xpos_in[worldid, pid] xquat = math.mul_quat(xquat_in[worldid, pid], body_quat_) for _ in range(jntnum): @@ -160,7 +156,6 @@ def _kinematics_level( xpos_out[worldid, bodyid] = xpos xquat_out[worldid, bodyid] = wp.normalize(xquat) - xmat_out[worldid, bodyid] = math.quat_to_mat(xquat) xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid % body_ipos.shape[0], bodyid], xquat) ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[worldid % body_iquat.shape[0], bodyid])) @@ -284,7 +279,7 @@ def kinematics(m: Model, d: Data): derived positions and orientations of geoms, sites, and flexible elements, based on the current joint positions and any attached mocap bodies. """ - wp.launch(_kinematics_root, dim=(d.nworld), inputs=[], outputs=[d.xpos, d.xquat, d.xmat, d.xipos, d.ximat]) + wp.launch(_kinematics_root, dim=(d.nworld), inputs=[], outputs=[d.xpos, d.xquat, d.xipos, d.ximat]) for i in range(1, len(m.body_tree)): body_tree = m.body_tree[i] @@ -310,10 +305,9 @@ def kinematics(m: Model, d: Data): d.mocap_quat, d.xpos, d.xquat, - d.xmat, body_tree, ], - outputs=[d.xpos, d.xquat, d.xmat, d.xipos, d.ximat, d.xanchor, d.xaxis], + outputs=[d.xpos, d.xquat, d.xipos, d.ximat, d.xanchor, d.xaxis], ) wp.launch( @@ -456,7 +450,7 @@ def _cdof( jnt_dofadr: wp.array(dtype=int), jnt_bodyid: wp.array(dtype=int), # Data in: - xmat_in: wp.array2d(dtype=wp.mat33), + xquat_in: wp.array2d(dtype=wp.quat), xanchor_in: wp.array2d(dtype=wp.vec3), xaxis_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3), @@ -468,7 +462,7 @@ def _cdof( dofid = jnt_dofadr[jntid] jnt_type_ = jnt_type[jntid] xaxis = xaxis_in[worldid, jntid] - xmat = wp.transpose(xmat_in[worldid, bodyid]) + xmat = wp.transpose(math.quat_to_mat(xquat_in[worldid, bodyid])) # compute com-anchor vector offset = subtree_com_in[worldid, body_rootid[bodyid]] - xanchor_in[worldid, jntid] @@ -522,7 +516,7 @@ def com_pos(m: Model, d: Data): wp.launch( _cdof, dim=(d.nworld, m.njnt), - inputs=[m.body_rootid, m.jnt_type, m.jnt_dofadr, m.jnt_bodyid, d.xmat, d.xanchor, d.xaxis, d.subtree_com], + inputs=[m.body_rootid, m.jnt_type, m.jnt_dofadr, m.jnt_bodyid, d.xquat, d.xanchor, d.xaxis, d.subtree_com], outputs=[d.cdof], ) @@ -1157,7 +1151,7 @@ def _cfrc_ext_equality( eq_data: wp.array2d(dtype=vec11), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), - xmat_in: wp.array2d(dtype=wp.mat33), + xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), efc_id_in: wp.array2d(dtype=int), efc_force_in: wp.array2d(dtype=float), @@ -1214,7 +1208,7 @@ def _cfrc_ext_equality( offset = site_pos[worldid, obj1] # transform point on body1: local -> global - pos = xmat_in[worldid, bodyid1] @ offset + xpos_in[worldid, bodyid1] + pos = math.rot_vec_quat(offset, xquat_in[worldid, bodyid1]) + xpos_in[worldid, bodyid1] # subtree CoM-based torque_force vector newpos = subtree_com_in[worldid, body_rootid[bodyid1]] @@ -1236,7 +1230,7 @@ def _cfrc_ext_equality( offset = site_pos[worldid, obj2] # transform point on body2: local -> global - pos = xmat_in[worldid, bodyid2] @ offset + xpos_in[worldid, bodyid2] + pos = math.rot_vec_quat(offset, xquat_in[worldid, bodyid2]) + xpos_in[worldid, bodyid2] # subtree CoM-based torque_force vector newpos = subtree_com_in[worldid, body_rootid[bodyid2]] @@ -1343,7 +1337,7 @@ def rne_postconstraint(m: Model, d: Data): m.eq_objtype, m.eq_data, d.xpos, - d.xmat, + d.xquat, d.subtree_com, d.efc.id, d.efc.force, diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index aa63d2a23..a3bcfbbca 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -120,7 +120,6 @@ def test_kinematics(self): d.xaxis, d.xpos, d.xquat, - d.xmat, d.xipos, d.ximat, d.geom_xpos, @@ -136,7 +135,6 @@ def test_kinematics(self): _assert_eq(d.xaxis.numpy()[0], mjd.xaxis, "xaxis") _assert_eq(d.xpos.numpy()[0], mjd.xpos, "xpos") _assert_eq(d.xquat.numpy()[0], mjd.xquat, "xquat") - _assert_eq(d.xmat.numpy()[0], mjd.xmat.reshape((-1, 3, 3)), "xmat") _assert_eq(d.xipos.numpy()[0], mjd.xipos, "xipos") _assert_eq(d.ximat.numpy()[0], mjd.ximat.reshape((-1, 3, 3)), "ximat") _assert_eq(d.geom_xpos.numpy()[0], mjd.geom_xpos, "geom_xpos") diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 294c118ee..b53e117b5 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1566,7 +1566,6 @@ class Data: sensordata: sensor data array (nworld, nsensordata,) xpos: Cartesian position of body frame (nworld, nbody, 3) xquat: Cartesian orientation of body frame (nworld, nbody, 4) - xmat: Cartesian orientation of body frame (nworld, nbody, 3, 3) xipos: Cartesian position of body com (nworld, nbody, 3) ximat: Cartesian orientation of body inertia (nworld, nbody, 3, 3) xanchor: Cartesian position of joint anchor (nworld, njnt, 3) @@ -1666,7 +1665,6 @@ class Data: sensordata: array("nworld", "nsensordata", float) xpos: array("nworld", "nbody", wp.vec3) xquat: array("nworld", "nbody", wp.quat) - xmat: array("nworld", "nbody", wp.mat33) xipos: array("nworld", "nbody", wp.vec3) ximat: array("nworld", "nbody", wp.mat33) xanchor: array("nworld", "njnt", wp.vec3) From b43507fc571e78b016688e9c6879e17db37daa2c Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Fri, 5 Dec 2025 17:47:29 +0100 Subject: [PATCH 002/107] fix tests Signed-off-by: Alain Denzler --- mujoco_warp/_src/forward_test.py | 3 +-- mujoco_warp/_src/io_test.py | 1 - 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/mujoco_warp/_src/forward_test.py b/mujoco_warp/_src/forward_test.py index 47797e92c..5eeb58e0b 100644 --- a/mujoco_warp/_src/forward_test.py +++ b/mujoco_warp/_src/forward_test.py @@ -315,7 +315,6 @@ def test_step1(self, xml, energy): step1_field = [ "xpos", "xquat", - "xmat", "xipos", "ximat", "xanchor", @@ -391,7 +390,7 @@ def _getattr(arr): d_arr, is_nefc = _getattr(arr) d_arr = d_arr.numpy()[0] mjd_arr = getattr(mjd, arr) - if arr in ["xmat", "ximat", "geom_xmat", "site_xmat", "cam_xmat"]: + if arr in ["ximat", "geom_xmat", "site_xmat", "cam_xmat"]: mjd_arr = mjd_arr.reshape(-1) d_arr = d_arr.reshape(-1) elif arr == "qM": diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index 29c975d86..78a94861d 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -126,7 +126,6 @@ def test_get_data_into(self, nworld, world_id): "act_dot", "xpos", "xquat", - "xmat", "xipos", "ximat", "xanchor", From 577367da78a58cd597cd0b114d73e3be472cb4d2 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:13:20 +0100 Subject: [PATCH 003/107] slow transition Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 5 +++-- mujoco_warp/_src/types.py | 4 ++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index f570e5529..a45d90805 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -47,13 +47,14 @@ def _kinematics_root( xquat_out: wp.array2d(dtype=wp.quat), xipos_out: wp.array2d(dtype=wp.vec3), ximat_out: wp.array2d(dtype=wp.mat33), + xiquat_out: wp.array2d(dtype=wp.quat), ): worldid = wp.tid() xpos_out[worldid, 0] = wp.vec3(0.0) xquat_out[worldid, 0] = wp.quat(1.0, 0.0, 0.0, 0.0) xipos_out[worldid, 0] = wp.vec3(0.0) ximat_out[worldid, 0] = wp.identity(n=3, dtype=wp.float32) - + xiquat_out[worldid, 0] = wp.quat(1.0, 0.0, 0.0, 0.0) @wp.kernel def _kinematics_level( @@ -279,7 +280,7 @@ def kinematics(m: Model, d: Data): derived positions and orientations of geoms, sites, and flexible elements, based on the current joint positions and any attached mocap bodies. """ - wp.launch(_kinematics_root, dim=(d.nworld), inputs=[], outputs=[d.xpos, d.xquat, d.xipos, d.ximat]) + wp.launch(_kinematics_root, dim=(d.nworld), inputs=[], outputs=[d.xpos, d.xquat, d.xipos, d.ximat, d.xiquat]) for i in range(1, len(m.body_tree)): body_tree = m.body_tree[i] diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index b53e117b5..258889536 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1641,6 +1641,7 @@ class Data: collision_pairid: ids from broadphase (naconmax, 2) collision_worldid: collision world ids from broadphase (naconmax,) ncollision: collision count from broadphase (1,) + xiquat: Cartesian orientation of body inertia (nworld, nbody, 4) """ solver_niter: array("nworld", int) @@ -1738,3 +1739,6 @@ class Data: collision_pairid: array("naconmax", wp.vec2i) collision_worldid: array("naconmax", int) ncollision: array(1, int) + + xiquat: array("nworld", "nbody", wp.quat) + From bc47238e1d1ace4b9de87478bbe97a6c49136561 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:16:49 +0100 Subject: [PATCH 004/107] set xiquat Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index a45d90805..c23cca8c2 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -85,6 +85,7 @@ def _kinematics_level( xquat_out: wp.array2d(dtype=wp.quat), xipos_out: wp.array2d(dtype=wp.vec3), ximat_out: wp.array2d(dtype=wp.mat33), + xiquat_out: wp.array2d(dtype=wp.quat), xanchor_out: wp.array2d(dtype=wp.vec3), xaxis_out: wp.array2d(dtype=wp.vec3), ): @@ -159,6 +160,7 @@ def _kinematics_level( xquat_out[worldid, bodyid] = wp.normalize(xquat) xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid % body_ipos.shape[0], bodyid], xquat) ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[worldid % body_iquat.shape[0], bodyid])) + xiquat_out[worldid, bodyid] = math.mul_quat(xquat, body_iquat[worldid % body_iquat.shape[0], bodyid]) @wp.kernel @@ -308,7 +310,7 @@ def kinematics(m: Model, d: Data): d.xquat, body_tree, ], - outputs=[d.xpos, d.xquat, d.xipos, d.ximat, d.xanchor, d.xaxis], + outputs=[d.xpos, d.xquat, d.xipos, d.ximat, d.xiquat, d.xanchor, d.xaxis], ) wp.launch( From fada45bc420388e0bfa0799c63e69a19fd83c0b4 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:23:49 +0100 Subject: [PATCH 005/107] first conversions Signed-off-by: Alain Denzler --- mujoco_warp/_src/passive.py | 9 ++++++--- mujoco_warp/_src/passive_test.py | 1 + 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 66d6317b9..8c1b32ff0 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -280,7 +280,7 @@ def _fluid_force( body_fluid_ellipsoid: wp.array(dtype=bool), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xmat_in: wp.array2d(dtype=wp.mat33), subtree_com_in: wp.array2d(dtype=wp.vec3), @@ -302,7 +302,7 @@ def _fluid_force( # Body kinematics xipos = xipos_in[worldid, bodyid] - rot = ximat_in[worldid, bodyid] + rot = math.quat_to_mat(xiquat_in[worldid, bodyid]) rotT = wp.transpose(rot) cvel = cvel_in[worldid, bodyid] ang_global = wp.spatial_top(cvel) @@ -488,6 +488,9 @@ def _fluid_force( def _fluid(m: Model, d: Data): fluid_applied = wp.empty((d.nworld, m.nbody), dtype=wp.spatial_vector) + print(d.ximat.numpy()[0]) + print(d.xiquat.numpy()[0]) + wp.launch( _fluid_force, dim=(d.nworld, m.nbody), @@ -505,7 +508,7 @@ def _fluid(m: Model, d: Data): m.geom_fluid, m.body_fluid_ellipsoid, d.xipos, - d.ximat, + d.xiquat, d.geom_xpos, d.geom_xmat, d.subtree_com, diff --git a/mujoco_warp/_src/passive_test.py b/mujoco_warp/_src/passive_test.py index 2bc6220de..5d11b6ec5 100644 --- a/mujoco_warp/_src/passive_test.py +++ b/mujoco_warp/_src/passive_test.py @@ -90,6 +90,7 @@ def test_fluid(self, density, viscosity, wind0, wind1, wind2): for arr in (d.qfrc_passive, d.qfrc_fluid): arr.zero_() + mjw.kinematics(m, d) mjw.passive(m, d) _assert_eq(d.qfrc_passive.numpy()[0], mjd.qfrc_passive, "qfrc_passive") From ccb951ebd74eb67bc5a6307d45b191780b6d9cbe Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:25:08 +0100 Subject: [PATCH 006/107] remove prints Signed-off-by: Alain Denzler --- mujoco_warp/_src/passive.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 8c1b32ff0..4c1a01f02 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -488,9 +488,6 @@ def _fluid_force( def _fluid(m: Model, d: Data): fluid_applied = wp.empty((d.nworld, m.nbody), dtype=wp.spatial_vector) - print(d.ximat.numpy()[0]) - print(d.xiquat.numpy()[0]) - wp.launch( _fluid_force, dim=(d.nworld, m.nbody), From 1f4d7b02c13cd977b7d8f0505af2718ad087c11e Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:29:55 +0100 Subject: [PATCH 007/107] conversions in sensor Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 18 +++++++++--------- mujoco_warp/_src/sensor_test.py | 1 + 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 44b08f529..6c66f626f 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -281,7 +281,7 @@ def _frame_pos( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xmat_in: wp.array2d(dtype=wp.mat33), site_xpos_in: wp.array2d(dtype=wp.vec3), @@ -313,7 +313,7 @@ def _frame_pos( if reftype == ObjType.BODY: xpos_ref = xipos_in[worldid, refid] - xmat_ref = ximat_in[worldid, refid] + xmat_ref = math.quat_to_mat(xiquat_in[worldid, refid]) elif objtype == ObjType.XBODY: xpos_ref = xpos_in[worldid, refid] xmat_ref = math.quat_to_mat(xquat_in[worldid, refid]) @@ -338,7 +338,7 @@ def _frame_pos( def _frame_axis( # Data in: xquat_in: wp.array2d(dtype=wp.quat), - ximat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), geom_xmat_in: wp.array2d(dtype=wp.mat33), site_xmat_in: wp.array2d(dtype=wp.mat33), cam_xmat_in: wp.array2d(dtype=wp.mat33), @@ -351,7 +351,7 @@ def _frame_axis( frame_axis: int, ) -> wp.vec3: if objtype == ObjType.BODY: - xmat = ximat_in[worldid, objid] + xmat = math.quat_to_mat(xiquat_in[worldid, objid]) axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) elif objtype == ObjType.XBODY: xmat = math.quat_to_mat(xquat_in[worldid, objid]) @@ -372,7 +372,7 @@ def _frame_axis( return axis if reftype == ObjType.BODY: - xmat_ref = ximat_in[worldid, refid] + xmat_ref = math.quat_to_mat(xiquat_in[worldid, refid]) elif reftype == ObjType.XBODY: xmat_ref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: @@ -492,7 +492,7 @@ def _sensor_pos( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xmat_in: wp.array2d(dtype=wp.mat33), site_xpos_in: wp.array2d(dtype=wp.vec3), @@ -546,7 +546,7 @@ def _sensor_pos( xpos_in, xquat_in, xipos_in, - ximat_in, + xiquat_in, geom_xpos_in, geom_xmat_in, site_xpos_in, @@ -571,7 +571,7 @@ def _sensor_pos( elif sensortype == SensorType.FRAMEZAXIS: axis = 2 vec3 = _frame_axis( - xquat_in, ximat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis + xquat_in, xiquat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEQUAT: @@ -855,7 +855,7 @@ def sensor_pos(m: Model, d: Data): d.xpos, d.xquat, d.xipos, - d.ximat, + d.xiquat, d.geom_xpos, d.geom_xmat, d.site_xpos, diff --git a/mujoco_warp/_src/sensor_test.py b/mujoco_warp/_src/sensor_test.py index 6b9eb4461..3e8c18eb4 100644 --- a/mujoco_warp/_src/sensor_test.py +++ b/mujoco_warp/_src/sensor_test.py @@ -292,6 +292,7 @@ def test_sensor(self): d.sensordata.zero_() + mjw.kinematics(m, d) mjw.sensor_pos(m, d) mjw.sensor_vel(m, d) mjw.sensor_acc(m, d) From 31e52f8944eaa8478273c4a67b161511bfc41363 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:34:09 +0100 Subject: [PATCH 008/107] sensor_vel Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 6c66f626f..2fb6bc5fc 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -1051,7 +1051,7 @@ def _frame_linvel( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xmat_in: wp.array2d(dtype=wp.mat33), site_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1082,7 +1082,7 @@ def _frame_linvel( if reftype == ObjType.BODY: xposref = xipos_in[worldid, refid] - xmatref = ximat_in[worldid, refid] + xmatref = math.quat_to_mat(xiquat_in[worldid, refid]) elif reftype == ObjType.XBODY: xposref = xpos_in[worldid, refid] xmatref = math.quat_to_mat(xquat_in[worldid, refid]) @@ -1157,7 +1157,7 @@ def _frame_angvel( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xmat_in: wp.array2d(dtype=wp.mat33), site_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1193,7 +1193,7 @@ def _frame_angvel( if refid > -1: if reftype == ObjType.BODY: - xmatref = ximat_in[worldid, refid] + xmatref = math.quat_to_mat(xiquat_in[worldid, refid]) elif reftype == ObjType.XBODY: xmatref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: @@ -1260,7 +1260,7 @@ def _sensor_vel( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xmat_in: wp.array2d(dtype=wp.mat33), site_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1312,7 +1312,7 @@ def _sensor_vel( xpos_in, xquat_in, xipos_in, - ximat_in, + xiquat_in, geom_xpos_in, geom_xmat_in, site_xpos_in, @@ -1340,7 +1340,7 @@ def _sensor_vel( xpos_in, xquat_in, xipos_in, - ximat_in, + xiquat_in, geom_xpos_in, geom_xmat_in, site_xpos_in, @@ -1395,7 +1395,7 @@ def sensor_vel(m: Model, d: Data): d.xpos, d.xquat, d.xipos, - d.ximat, + d.xiquat, d.geom_xpos, d.geom_xmat, d.site_xpos, From 8e1e8256058dc186a784f6a7a6a93d2ddd59d113 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:38:08 +0100 Subject: [PATCH 009/107] cinert Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 6 +++--- mujoco_warp/_src/smooth_test.py | 1 + 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index c23cca8c2..7b5fcc071 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -407,13 +407,13 @@ def _cinert( body_inertia: wp.array2d(dtype=wp.vec3), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), # Data out: cinert_out: wp.array2d(dtype=vec10), ): worldid, bodyid = wp.tid() - mat = ximat_in[worldid, bodyid] + mat = math.quat_to_mat(xiquat_in[worldid, bodyid]) inert = body_inertia[worldid % body_inertia.shape[0], bodyid] mass = body_mass[worldid % body_mass.shape[0], bodyid] dif = xipos_in[worldid, bodyid] - subtree_com_in[worldid, body_rootid[bodyid]] @@ -513,7 +513,7 @@ def com_pos(m: Model, d: Data): wp.launch( _cinert, dim=(d.nworld, m.nbody), - inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.ximat, d.subtree_com], + inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.xiquat, d.subtree_com], outputs=[d.cinert], ) wp.launch( diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index a3bcfbbca..276642a60 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -149,6 +149,7 @@ def test_com_pos(self): for arr in (d.subtree_com, d.cinert, d.cdof): arr.zero_() + mjw.kinematics(m, d) mjw.com_pos(m, d) _assert_eq(d.subtree_com.numpy()[0], mjd.subtree_com, "subtree_com") _assert_eq(d.cinert.numpy()[0], mjd.cinert, "cinert") From 37d95876d51f43b5dbecaf17486379a990a9db03 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:40:08 +0100 Subject: [PATCH 010/107] subtree_vel_forward Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 7b5fcc071..51737e22c 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -2447,7 +2447,7 @@ def _subtree_vel_forward( body_inertia: wp.array2d(dtype=wp.vec3), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), - ximat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), # Data out: @@ -2463,7 +2463,7 @@ def _subtree_vel_forward( ang = wp.spatial_top(cvel) lin = wp.spatial_bottom(cvel) xipos = xipos_in[worldid, bodyid] - ximat = ximat_in[worldid, bodyid] + ximat = math.quat_to_mat(xiquat_in[worldid, bodyid]) subtree_com_root = subtree_com_in[worldid, body_rootid[bodyid]] # update linear velocity @@ -2561,7 +2561,7 @@ def subtree_vel(m: Model, d: Data): wp.launch( _subtree_vel_forward, dim=(d.nworld, m.nbody), - inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.ximat, d.subtree_com, d.cvel], + inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.xiquat, d.subtree_com, d.cvel], outputs=[d.subtree_linvel, d.subtree_angmom, d.subtree_bodyvel], ) From 10ac877818b0f9fab143ed355762077d138f065f Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:49:34 +0100 Subject: [PATCH 011/107] no more ximat Signed-off-by: Alain Denzler --- mujoco_warp/_src/forward_test.py | 1 - mujoco_warp/_src/smooth.py | 8 ++------ mujoco_warp/_src/smooth_test.py | 6 +++++- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/forward_test.py b/mujoco_warp/_src/forward_test.py index 5eeb58e0b..11d782478 100644 --- a/mujoco_warp/_src/forward_test.py +++ b/mujoco_warp/_src/forward_test.py @@ -316,7 +316,6 @@ def test_step1(self, xml, energy): "xpos", "xquat", "xipos", - "ximat", "xanchor", "xaxis", "geom_xpos", diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 51737e22c..cced79bfc 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -46,14 +46,12 @@ def _kinematics_root( xpos_out: wp.array2d(dtype=wp.vec3), xquat_out: wp.array2d(dtype=wp.quat), xipos_out: wp.array2d(dtype=wp.vec3), - ximat_out: wp.array2d(dtype=wp.mat33), xiquat_out: wp.array2d(dtype=wp.quat), ): worldid = wp.tid() xpos_out[worldid, 0] = wp.vec3(0.0) xquat_out[worldid, 0] = wp.quat(1.0, 0.0, 0.0, 0.0) xipos_out[worldid, 0] = wp.vec3(0.0) - ximat_out[worldid, 0] = wp.identity(n=3, dtype=wp.float32) xiquat_out[worldid, 0] = wp.quat(1.0, 0.0, 0.0, 0.0) @wp.kernel @@ -84,7 +82,6 @@ def _kinematics_level( xpos_out: wp.array2d(dtype=wp.vec3), xquat_out: wp.array2d(dtype=wp.quat), xipos_out: wp.array2d(dtype=wp.vec3), - ximat_out: wp.array2d(dtype=wp.mat33), xiquat_out: wp.array2d(dtype=wp.quat), xanchor_out: wp.array2d(dtype=wp.vec3), xaxis_out: wp.array2d(dtype=wp.vec3), @@ -159,7 +156,6 @@ def _kinematics_level( xpos_out[worldid, bodyid] = xpos xquat_out[worldid, bodyid] = wp.normalize(xquat) xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid % body_ipos.shape[0], bodyid], xquat) - ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[worldid % body_iquat.shape[0], bodyid])) xiquat_out[worldid, bodyid] = math.mul_quat(xquat, body_iquat[worldid % body_iquat.shape[0], bodyid]) @@ -282,7 +278,7 @@ def kinematics(m: Model, d: Data): derived positions and orientations of geoms, sites, and flexible elements, based on the current joint positions and any attached mocap bodies. """ - wp.launch(_kinematics_root, dim=(d.nworld), inputs=[], outputs=[d.xpos, d.xquat, d.xipos, d.ximat, d.xiquat]) + wp.launch(_kinematics_root, dim=(d.nworld), inputs=[], outputs=[d.xpos, d.xquat, d.xipos, d.xiquat]) for i in range(1, len(m.body_tree)): body_tree = m.body_tree[i] @@ -310,7 +306,7 @@ def kinematics(m: Model, d: Data): d.xquat, body_tree, ], - outputs=[d.xpos, d.xquat, d.xipos, d.ximat, d.xiquat, d.xanchor, d.xaxis], + outputs=[d.xpos, d.xquat, d.xipos, d.xiquat, d.xanchor, d.xaxis], ) wp.launch( diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index 276642a60..c8e9e5546 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -121,7 +121,7 @@ def test_kinematics(self): d.xpos, d.xquat, d.xipos, - d.ximat, + d.xiquat, d.geom_xpos, d.geom_xmat, d.site_xpos, @@ -131,6 +131,10 @@ def test_kinematics(self): mjw.kinematics(m, d) + ximat = np.zeros((m.nbody, 9), dtype=np.float64) + for i in range(m.nbody): + mujoco.mju_quat2Mat(ximat[i], d.xiquat.numpy()[0, i]) + _assert_eq(d.xanchor.numpy()[0], mjd.xanchor, "xanchor") _assert_eq(d.xaxis.numpy()[0], mjd.xaxis, "xaxis") _assert_eq(d.xpos.numpy()[0], mjd.xpos, "xpos") From 2c14b6ded7d4bb71d4b041102d9dad4dc27cf68e Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:52:05 +0100 Subject: [PATCH 012/107] properly test xiquat Signed-off-by: Alain Denzler --- mujoco_warp/_src/forward_test.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/mujoco_warp/_src/forward_test.py b/mujoco_warp/_src/forward_test.py index 11d782478..c402b36af 100644 --- a/mujoco_warp/_src/forward_test.py +++ b/mujoco_warp/_src/forward_test.py @@ -429,6 +429,13 @@ def _getattr(arr): _assert_eq(d_arr, mjd_arr, arr) + # compare xiquat (warp) to ximat (mujoco) by converting quaternion to matrix + xiquat = d.xiquat.numpy()[0] + ximat = np.zeros((xiquat.shape[0], 9)) + for i in range(xiquat.shape[0]): + mujoco.mju_quat2Mat(ximat[i], xiquat[i]) + _assert_eq(ximat.reshape(-1), mjd.ximat.reshape(-1), "xiquat->ximat") + # TODO(team): sensor_pos # TODO(team): sensor_vel From 1db0e22a9e9183ab7708042b5770d27bcbb40e4a Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 14:57:07 +0100 Subject: [PATCH 013/107] proper put_data conversion Signed-off-by: Alain Denzler --- mujoco_warp/_src/io.py | 7 +++++++ mujoco_warp/_src/passive_test.py | 1 - mujoco_warp/_src/sensor_test.py | 1 - mujoco_warp/_src/smooth_test.py | 1 - 4 files changed, 7 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index fcd8803b3..da8126309 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -781,6 +781,7 @@ def put_data( "ne_ten": None, "ne_flex": None, "nsolving": None, + "xiquat": None, } for f in dataclasses.fields(types.Data): if f.name in d_kwargs: @@ -826,6 +827,12 @@ def put_data( mujoco.mju_sparse2dense(actuator_moment, mjd.actuator_moment, mjd.moment_rownnz, mjd.moment_rowadr, mjd.moment_colind) d.actuator_moment = wp.array(np.full((nworld, mjm.nu, mjm.nv), actuator_moment), dtype=float) + # convert ximat to xiquat + xiquat = np.zeros((mjm.nbody, 4)) + for i in range(mjm.nbody): + mujoco.mju_mat2Quat(xiquat[i], mjd.ximat[i]) + d.xiquat = wp.array(np.full((nworld, mjm.nbody, 4), xiquat), dtype=wp.quat) + d.nacon = wp.array([mjd.ncon * nworld], dtype=int) d.ne_connect = wp.full(nworld, 3 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_CONNECT) & mjd.eq_active), dtype=int) d.ne_weld = wp.full(nworld, 6 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_WELD) & mjd.eq_active), dtype=int) diff --git a/mujoco_warp/_src/passive_test.py b/mujoco_warp/_src/passive_test.py index 5d11b6ec5..2bc6220de 100644 --- a/mujoco_warp/_src/passive_test.py +++ b/mujoco_warp/_src/passive_test.py @@ -90,7 +90,6 @@ def test_fluid(self, density, viscosity, wind0, wind1, wind2): for arr in (d.qfrc_passive, d.qfrc_fluid): arr.zero_() - mjw.kinematics(m, d) mjw.passive(m, d) _assert_eq(d.qfrc_passive.numpy()[0], mjd.qfrc_passive, "qfrc_passive") diff --git a/mujoco_warp/_src/sensor_test.py b/mujoco_warp/_src/sensor_test.py index 3e8c18eb4..6b9eb4461 100644 --- a/mujoco_warp/_src/sensor_test.py +++ b/mujoco_warp/_src/sensor_test.py @@ -292,7 +292,6 @@ def test_sensor(self): d.sensordata.zero_() - mjw.kinematics(m, d) mjw.sensor_pos(m, d) mjw.sensor_vel(m, d) mjw.sensor_acc(m, d) diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index c8e9e5546..3a03ae0aa 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -153,7 +153,6 @@ def test_com_pos(self): for arr in (d.subtree_com, d.cinert, d.cdof): arr.zero_() - mjw.kinematics(m, d) mjw.com_pos(m, d) _assert_eq(d.subtree_com.numpy()[0], mjd.subtree_com, "subtree_com") _assert_eq(d.cinert.numpy()[0], mjd.cinert, "cinert") From b1e359b5604e0c02b6f0392c1f32c0584c1103a3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 15:03:54 +0100 Subject: [PATCH 014/107] better conversions Signed-off-by: Alain Denzler --- mujoco_warp/_src/forward_test.py | 2 +- mujoco_warp/_src/io.py | 7 ++++++- mujoco_warp/_src/io_test.py | 2 +- mujoco_warp/_src/smooth_test.py | 2 +- mujoco_warp/_src/types.py | 2 -- 5 files changed, 9 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/forward_test.py b/mujoco_warp/_src/forward_test.py index c402b36af..3d00d414e 100644 --- a/mujoco_warp/_src/forward_test.py +++ b/mujoco_warp/_src/forward_test.py @@ -389,7 +389,7 @@ def _getattr(arr): d_arr, is_nefc = _getattr(arr) d_arr = d_arr.numpy()[0] mjd_arr = getattr(mjd, arr) - if arr in ["ximat", "geom_xmat", "site_xmat", "cam_xmat"]: + if arr in ["geom_xmat", "site_xmat", "cam_xmat"]: mjd_arr = mjd_arr.reshape(-1) d_arr = d_arr.reshape(-1) elif arr == "qM": diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index da8126309..328644e05 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -902,6 +902,12 @@ def get_data_into( efc_idx = efc_idx[:nefc] # dont emit indices for overflow constraints + # convert xiquat to ximat + ximat = np.zeros((mjm.nbody, 9)) + for i in range(mjm.nbody): + mujoco.mju_quat2Mat(ximat[i], d.xiquat.numpy()[world_id, i]) + result.ximat[:] = ximat.reshape((-1, 9)) + result.solver_niter[0] = d.solver_niter.numpy()[world_id] result.ncon = ncon result.ne = ne @@ -938,7 +944,6 @@ def get_data_into( xmat[:, 8] = q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2 + q[:, 3] ** 2 result.xmat[:] = xmat result.xipos[:] = d.xipos.numpy()[world_id] - result.ximat[:] = d.ximat.numpy()[world_id].reshape((-1, 9)) result.xanchor[:] = d.xanchor.numpy()[world_id] result.xaxis[:] = d.xaxis.numpy()[world_id] result.geom_xpos[:] = d.geom_xpos.numpy()[world_id] diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index 78a94861d..a6994e972 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -127,7 +127,7 @@ def test_get_data_into(self, nworld, world_id): "xpos", "xquat", "xipos", - "ximat", + #"ximat", "xanchor", "xaxis", "geom_xpos", diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index 3a03ae0aa..35c0d062f 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -140,7 +140,7 @@ def test_kinematics(self): _assert_eq(d.xpos.numpy()[0], mjd.xpos, "xpos") _assert_eq(d.xquat.numpy()[0], mjd.xquat, "xquat") _assert_eq(d.xipos.numpy()[0], mjd.xipos, "xipos") - _assert_eq(d.ximat.numpy()[0], mjd.ximat.reshape((-1, 3, 3)), "ximat") + _assert_eq(ximat, mjd.ximat, "ximat") _assert_eq(d.geom_xpos.numpy()[0], mjd.geom_xpos, "geom_xpos") _assert_eq(d.geom_xmat.numpy()[0], mjd.geom_xmat.reshape((-1, 3, 3)), "geom_xmat") _assert_eq(d.site_xpos.numpy()[0], mjd.site_xpos, "site_xpos") diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 258889536..f3dff403c 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1567,7 +1567,6 @@ class Data: xpos: Cartesian position of body frame (nworld, nbody, 3) xquat: Cartesian orientation of body frame (nworld, nbody, 4) xipos: Cartesian position of body com (nworld, nbody, 3) - ximat: Cartesian orientation of body inertia (nworld, nbody, 3, 3) xanchor: Cartesian position of joint anchor (nworld, njnt, 3) xaxis: Cartesian joint axis (nworld, njnt, 3) geom_xpos: Cartesian geom position (nworld, ngeom, 3) @@ -1667,7 +1666,6 @@ class Data: xpos: array("nworld", "nbody", wp.vec3) xquat: array("nworld", "nbody", wp.quat) xipos: array("nworld", "nbody", wp.vec3) - ximat: array("nworld", "nbody", wp.mat33) xanchor: array("nworld", "njnt", wp.vec3) xaxis: array("nworld", "njnt", wp.vec3) geom_xpos: array("nworld", "ngeom", wp.vec3) From bb8eaf15d7523384b8ec2a76fe0d332a74aceada Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 15:06:58 +0100 Subject: [PATCH 015/107] initial infra Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 2 ++ mujoco_warp/_src/types.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index cced79bfc..a88546b9f 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -174,6 +174,7 @@ def _geom_local_to_global( # Data out: geom_xpos_out: wp.array2d(dtype=wp.vec3), geom_xmat_out: wp.array2d(dtype=wp.mat33), + geom_xquat_out: wp.array2d(dtype=wp.quat), ): worldid, geomid = wp.tid() bodyid = geom_bodyid[geomid] @@ -187,6 +188,7 @@ def _geom_local_to_global( xquat = xquat_in[worldid, bodyid] geom_xpos_out[worldid, geomid] = xpos + math.rot_vec_quat(geom_pos[worldid % geom_pos.shape[0], geomid], xquat) geom_xmat_out[worldid, geomid] = math.quat_to_mat(math.mul_quat(xquat, geom_quat[worldid % geom_quat.shape[0], geomid])) + geom_xquat_out[worldid, geomid] = xquat @wp.kernel diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index f3dff403c..7c36559fc 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1641,6 +1641,7 @@ class Data: collision_worldid: collision world ids from broadphase (naconmax,) ncollision: collision count from broadphase (1,) xiquat: Cartesian orientation of body inertia (nworld, nbody, 4) + geom_xquat: Cartesian orientation of geom (nworld, ngeom, 4) """ solver_niter: array("nworld", int) @@ -1739,4 +1740,5 @@ class Data: ncollision: array(1, int) xiquat: array("nworld", "nbody", wp.quat) + geom_xquat: array("nworld", "ngeom", wp.quat) From b3768ef13611702ab57f8eae8ec472d955e3afdd Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 15:08:31 +0100 Subject: [PATCH 016/107] kernel call needs it as well Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index a88546b9f..1686afef9 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -315,7 +315,7 @@ def kinematics(m: Model, d: Data): _geom_local_to_global, dim=(d.nworld, m.ngeom), inputs=[m.body_rootid, m.body_weldid, m.body_mocapid, m.geom_bodyid, m.geom_pos, m.geom_quat, d.xpos, d.xquat], - outputs=[d.geom_xpos, d.geom_xmat], + outputs=[d.geom_xpos, d.geom_xmat, d.geom_xquat], ) wp.launch( From 243b8bfed814057f58627d83809c4b57f5d263d2 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 15:09:56 +0100 Subject: [PATCH 017/107] spatial_tendon_geom Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 1686afef9..4f899da34 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -2694,7 +2694,7 @@ def _spatial_geom_tendon( wrap_pulley_scale: wp.array(dtype=float), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3), cdof_in: wp.array2d(dtype=wp.spatial_vector), @@ -2722,7 +2722,7 @@ def _spatial_geom_tendon( # get geom information geom_xpos = geom_xpos_in[worldid, wrap_objid_geom] - geom_xmat = geom_xmat_in[worldid, wrap_objid_geom] + geom_xmat = math.quat_to_mat(geom_xquat_in[worldid, wrap_objid_geom]) geomsize = geom_size[worldid, wrap_objid_geom][0] geom_type = wrap_type[wrap_adr] @@ -3052,7 +3052,7 @@ def tendon(m: Model, d: Data): m.wrap_geom_adr, m.wrap_pulley_scale, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.site_xpos, d.subtree_com, d.cdof, From 7ab5218fc127d4c76ebe02c80d218212cb09cd05 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 15:16:19 +0100 Subject: [PATCH 018/107] better put_data Signed-off-by: Alain Denzler --- mujoco_warp/_src/io.py | 12 +++++++----- mujoco_warp/_src/passive.py | 6 +++--- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 328644e05..723a29417 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -782,6 +782,7 @@ def put_data( "ne_flex": None, "nsolving": None, "xiquat": None, + "geom_xquat": None, } for f in dataclasses.fields(types.Data): if f.name in d_kwargs: @@ -827,11 +828,12 @@ def put_data( mujoco.mju_sparse2dense(actuator_moment, mjd.actuator_moment, mjd.moment_rownnz, mjd.moment_rowadr, mjd.moment_colind) d.actuator_moment = wp.array(np.full((nworld, mjm.nu, mjm.nv), actuator_moment), dtype=float) - # convert ximat to xiquat - xiquat = np.zeros((mjm.nbody, 4)) - for i in range(mjm.nbody): - mujoco.mju_mat2Quat(xiquat[i], mjd.ximat[i]) - d.xiquat = wp.array(np.full((nworld, mjm.nbody, 4), xiquat), dtype=wp.quat) + # convert xmat fields to xquat + for xmat, attr, n in [(mjd.ximat, "xiquat", mjm.nbody), (mjd.geom_xmat, "geom_xquat", mjm.ngeom)]: + xquat = np.zeros((n, 4)) + for i in range(n): + mujoco.mju_mat2Quat(xquat[i], xmat[i]) + setattr(d, attr, wp.array(np.full((nworld, n, 4), xquat), dtype=wp.quat)) d.nacon = wp.array([mjd.ncon * nworld], dtype=int) d.ne_connect = wp.full(nworld, 3 * np.sum((mjm.eq_type == mujoco.mjtEq.mjEQ_CONNECT) & mjd.eq_active), dtype=int) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 4c1a01f02..1896f956f 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -282,7 +282,7 @@ def _fluid_force( xipos_in: wp.array2d(dtype=wp.vec3), xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), # Out: @@ -325,7 +325,7 @@ def _fluid_force( size = geom_size[worldid % geom_size.shape[0], geomid] semiaxes = _geom_semiaxes(size, geom_type[geomid]) - geom_rot = geom_xmat_in[worldid, geomid] + geom_rot = math.quat_to_mat(geom_xquat_in[worldid, geomid]) geom_rotT = wp.transpose(geom_rot) geom_pos = geom_xpos_in[worldid, geomid] @@ -507,7 +507,7 @@ def _fluid(m: Model, d: Data): d.xipos, d.xiquat, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.subtree_com, d.cvel, ], From de2341106d8ada2712ead7b32bd3c6b933516517 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 16:26:48 +0100 Subject: [PATCH 019/107] better smooth_test Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth_test.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index 35c0d062f..7ee5a4bb8 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -123,7 +123,7 @@ def test_kinematics(self): d.xipos, d.xiquat, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.site_xpos, d.site_xmat, ): @@ -135,6 +135,10 @@ def test_kinematics(self): for i in range(m.nbody): mujoco.mju_quat2Mat(ximat[i], d.xiquat.numpy()[0, i]) + geom_xmat = np.zeros((m.ngeom, 9), dtype=np.float64) + for i in range(m.ngeom): + mujoco.mju_quat2Mat(geom_xmat[i], d.geom_xquat.numpy()[0, i]) + _assert_eq(d.xanchor.numpy()[0], mjd.xanchor, "xanchor") _assert_eq(d.xaxis.numpy()[0], mjd.xaxis, "xaxis") _assert_eq(d.xpos.numpy()[0], mjd.xpos, "xpos") @@ -142,7 +146,7 @@ def test_kinematics(self): _assert_eq(d.xipos.numpy()[0], mjd.xipos, "xipos") _assert_eq(ximat, mjd.ximat, "ximat") _assert_eq(d.geom_xpos.numpy()[0], mjd.geom_xpos, "geom_xpos") - _assert_eq(d.geom_xmat.numpy()[0], mjd.geom_xmat.reshape((-1, 3, 3)), "geom_xmat") + _assert_eq(geom_xmat, mjd.geom_xmat, "geom_xmat") _assert_eq(d.site_xpos.numpy()[0], mjd.site_xpos, "site_xpos") _assert_eq(d.site_xmat.numpy()[0], mjd.site_xmat.reshape((-1, 3, 3)), "site_xmat") From 6f7df83fafd616e7a576fffab1ff0cc34a6ae8a3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 16:30:30 +0100 Subject: [PATCH 020/107] sensor_tactile Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 2fb6bc5fc..310358d4a 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -2110,7 +2110,7 @@ def _sensor_tactile( taxel_sensorid: wp.array(dtype=int), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), contact_geom_in: wp.array(dtype=wp.vec2i), @@ -2149,14 +2149,15 @@ def _sensor_tactile( # vertex local position vertid = taxel_vertadr[taxelid] - mesh_vertadr[mesh_id] pos = mesh_vert[vertid + mesh_vertadr[mesh_id]] + quat = geom_xquat_in[worldid, geom_id] # position in global frame - xpos = geom_xmat_in[worldid, geom_id] @ pos + xpos = math.rot_vec_quat(pos, quat) xpos += geom_xpos_in[worldid, geom_id] # position in other geom frame tmp = xpos - geom_xpos_in[worldid, geom] - lpos = wp.transpose(geom_xmat_in[worldid, geom]) @ tmp + lpos = wp.transpose(math.quat_to_mat(quat)) @ tmp plugin_id = geom_plugin_index[geom] @@ -2451,7 +2452,7 @@ def sensor_acc(m: Model, d: Data): m.taxel_vertadr, m.taxel_sensorid, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.subtree_com, d.cvel, d.contact.geom, From 1c50e3f6925b6a78bd9912ea61c2d48f0cb498ab Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 16:32:39 +0100 Subject: [PATCH 021/107] sensor_vel Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 310358d4a..69b24f6bd 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -1053,7 +1053,7 @@ def _frame_linvel( xipos_in: wp.array2d(dtype=wp.vec3), xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), site_xmat_in: wp.array2d(dtype=wp.mat33), cam_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1088,7 +1088,7 @@ def _frame_linvel( xmatref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: xposref = geom_xpos_in[worldid, refid] - xmatref = geom_xmat_in[worldid, refid] + xmatref = math.quat_to_mat(geom_xquat_in[worldid, refid]) elif reftype == ObjType.SITE: xposref = site_xpos_in[worldid, refid] xmatref = site_xmat_in[worldid, refid] @@ -1159,7 +1159,7 @@ def _frame_angvel( xipos_in: wp.array2d(dtype=wp.vec3), xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), site_xmat_in: wp.array2d(dtype=wp.mat33), cam_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1197,7 +1197,7 @@ def _frame_angvel( elif reftype == ObjType.XBODY: xmatref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: - xmatref = geom_xmat_in[worldid, refid] + xmatref = math.quat_to_mat(geom_xquat_in[worldid, refid]) elif reftype == ObjType.SITE: xmatref = site_xmat_in[worldid, refid] elif reftype == ObjType.CAMERA: @@ -1262,7 +1262,7 @@ def _sensor_vel( xipos_in: wp.array2d(dtype=wp.vec3), xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), site_xmat_in: wp.array2d(dtype=wp.mat33), cam_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1314,7 +1314,7 @@ def _sensor_vel( xipos_in, xiquat_in, geom_xpos_in, - geom_xmat_in, + geom_xquat_in, site_xpos_in, site_xmat_in, cam_xpos_in, @@ -1342,7 +1342,7 @@ def _sensor_vel( xipos_in, xiquat_in, geom_xpos_in, - geom_xmat_in, + geom_xquat_in, site_xpos_in, site_xmat_in, cam_xpos_in, @@ -1397,7 +1397,7 @@ def sensor_vel(m: Model, d: Data): d.xipos, d.xiquat, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.site_xpos, d.site_xmat, d.cam_xpos, From 11076fd778565dfc8461a65bb9441cab04eb943b Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 16:35:29 +0100 Subject: [PATCH 022/107] sensor_pos Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 69b24f6bd..b79725a89 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -283,7 +283,7 @@ def _frame_pos( xipos_in: wp.array2d(dtype=wp.vec3), xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), site_xmat_in: wp.array2d(dtype=wp.mat33), cam_xpos_in: wp.array2d(dtype=wp.vec3), @@ -319,7 +319,7 @@ def _frame_pos( xmat_ref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: xpos_ref = geom_xpos_in[worldid, refid] - xmat_ref = geom_xmat_in[worldid, refid] + xmat_ref = math.quat_to_mat(geom_xquat_in[worldid, refid]) elif reftype == ObjType.SITE: xpos_ref = site_xpos_in[worldid, refid] xmat_ref = site_xmat_in[worldid, refid] @@ -339,7 +339,7 @@ def _frame_axis( # Data in: xquat_in: wp.array2d(dtype=wp.quat), xiquat_in: wp.array2d(dtype=wp.quat), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), site_xmat_in: wp.array2d(dtype=wp.mat33), cam_xmat_in: wp.array2d(dtype=wp.mat33), # In: @@ -357,7 +357,7 @@ def _frame_axis( xmat = math.quat_to_mat(xquat_in[worldid, objid]) axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) elif objtype == ObjType.GEOM: - xmat = geom_xmat_in[worldid, objid] + xmat = math.quat_to_mat(geom_xquat_in[worldid, objid]) axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) elif objtype == ObjType.SITE: xmat = site_xmat_in[worldid, objid] @@ -376,7 +376,7 @@ def _frame_axis( elif reftype == ObjType.XBODY: xmat_ref = math.quat_to_mat(xquat_in[worldid, refid]) elif reftype == ObjType.GEOM: - xmat_ref = geom_xmat_in[worldid, refid] + xmat_ref = math.quat_to_mat(geom_xquat_in[worldid, refid]) elif reftype == ObjType.SITE: xmat_ref = site_xmat_in[worldid, refid] elif reftype == ObjType.CAMERA: @@ -494,7 +494,7 @@ def _sensor_pos( xipos_in: wp.array2d(dtype=wp.vec3), xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), site_xmat_in: wp.array2d(dtype=wp.mat33), cam_xpos_in: wp.array2d(dtype=wp.vec3), @@ -548,7 +548,7 @@ def _sensor_pos( xipos_in, xiquat_in, geom_xpos_in, - geom_xmat_in, + geom_xquat_in, site_xpos_in, site_xmat_in, cam_xpos_in, @@ -571,7 +571,7 @@ def _sensor_pos( elif sensortype == SensorType.FRAMEZAXIS: axis = 2 vec3 = _frame_axis( - xquat_in, xiquat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis + xquat_in, xiquat_in, geom_xquat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEQUAT: @@ -857,7 +857,7 @@ def sensor_pos(m: Model, d: Data): d.xipos, d.xiquat, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.site_xpos, d.site_xmat, d.cam_xpos, From 3f4757923e6c8d434ca86a18d5885e8efbdeedf2 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 16:57:59 +0100 Subject: [PATCH 023/107] crashing tests Signed-off-by: Alain Denzler --- mujoco_warp/_src/broadphase_test.py | 10 ++++++---- mujoco_warp/_src/collision_convex.py | 15 ++++++++------- mujoco_warp/_src/collision_driver.py | 17 +++++++++-------- mujoco_warp/_src/collision_gjk_test.py | 10 ++++++---- mujoco_warp/_src/collision_primitive.py | 13 +++++++------ mujoco_warp/_src/collision_sdf.py | 6 +++--- mujoco_warp/_src/io.py | 4 ++++ mujoco_warp/_src/ray.py | 11 ++++++----- 8 files changed, 49 insertions(+), 37 deletions(-) diff --git a/mujoco_warp/_src/broadphase_test.py b/mujoco_warp/_src/broadphase_test.py index 1e1e35c63..30939c29d 100644 --- a/mujoco_warp/_src/broadphase_test.py +++ b/mujoco_warp/_src/broadphase_test.py @@ -15,6 +15,7 @@ """Tests for broadphase functions.""" +import mujoco import numpy as np import warp as wp from absl.testing import absltest @@ -145,10 +146,11 @@ def test_broadphase(self, broadphase, filter): np.vstack([np.expand_dims(mjd1.geom_xpos, axis=0), np.expand_dims(mjd2.geom_xpos, axis=0)]), dtype=wp.vec3, ) - d3.geom_xmat = wp.array( - np.vstack([np.expand_dims(mjd1.geom_xmat, axis=0), np.expand_dims(mjd2.geom_xmat, axis=0)]), - dtype=wp.mat33, - ) + geom_xquat = np.zeros((2, mjm.ngeom, 4)) + for j, mjd in enumerate([mjd1, mjd2]): + for i in range(mjm.ngeom): + mujoco.mju_mat2Quat(geom_xquat[j, i], mjd.geom_xmat[i]) + d3.geom_xquat = wp.array(geom_xquat, dtype=wp.quat) broadphase_caller(m, d3) ncollision = d3.ncollision.numpy()[0] diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index b626411ec..3f7e97813 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -25,6 +25,7 @@ from .collision_primitive import write_contact from .math import make_frame from .math import upper_trid_index +from .math import quat_to_mat from .types import MJ_MAX_EPAFACES from .types import MJ_MAX_EPAHORIZON from .types import MJ_MAXCONPAIR @@ -92,7 +93,7 @@ def _hfield_filter( hfield_size: wp.array(dtype=wp.vec4), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, g1: int, @@ -111,7 +112,7 @@ def _hfield_filter( margin_id = worldid % geom_margin.shape[0] pos1 = geom_xpos_in[worldid, g1] - mat1 = geom_xmat_in[worldid, g1] + mat1 = quat_to_mat(geom_xquat_in[worldid, g1]) mat1T = wp.transpose(mat1) pos2 = geom_xpos_in[worldid, g2] pos = mat1T @ (pos2 - pos1) @@ -132,7 +133,7 @@ def _hfield_filter( if -size1[3] > pos[2] + r2 + margin: # down return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf - mat2 = geom_xmat_in[worldid, g2] + mat2 = quat_to_mat(geom_xquat_in[worldid, g2]) mat = mat1T @ mat2 # aabb for geom in height field frame @@ -426,7 +427,7 @@ def ccd_kernel( # Data in: naconmax_in: int, geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), collision_pair_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array(dtype=wp.vec2i), collision_worldid_in: wp.array(dtype=int), @@ -486,7 +487,7 @@ def ccd_kernel( # height field filter if wp.static(is_hfield): no_hf_collision, xmin, xmax, ymin, ymax, zmin, zmax = _hfield_filter( - geom_dataid, geom_aabb, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xmat_in, worldid, g1, g2 + geom_dataid, geom_aabb, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xquat_in, worldid, g1, g2 ) if no_hf_collision: return @@ -532,7 +533,7 @@ def ccd_kernel( mesh_polymapnum, mesh_polymap, geom_xpos_in, - geom_xmat_in, + geom_xquat_in, geoms, worldid, ) @@ -1070,7 +1071,7 @@ def convex_narrowphase(m: Model, d: Data): m.pair_friction, d.naconmax, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.collision_pair, d.collision_pairid, d.collision_worldid, diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index a9cd39c97..c9a2ee67c 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -21,6 +21,7 @@ from .collision_primitive import primitive_narrowphase from .collision_sdf import sdf_narrowphase from .math import upper_tri_index +from .math import quat_to_mat from .types import MJ_MAXVAL from .types import BroadphaseFilter from .types import BroadphaseType @@ -239,7 +240,7 @@ def func( geom_margin: wp.array2d(dtype=float), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), # In: geom1: int, geom2: int, @@ -259,7 +260,7 @@ def func( margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0 margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2] xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2] - xmat1, xmat2 = geom_xmat_in[worldid, geom1], geom_xmat_in[worldid, geom2] + xmat1, xmat2 = quat_to_mat(geom_xquat_in[worldid, geom1]), quat_to_mat(geom_xquat_in[worldid, geom2]) if rbound1 == 0.0 or rbound2 == 0.0: if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE): @@ -415,7 +416,7 @@ def kernel( nworld_in: int, naconmax_in: int, geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), # In: sort_index_in: wp.array2d(dtype=int), cumulative_sum_in: wp.array(dtype=int), @@ -460,7 +461,7 @@ def kernel( if ( wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))( - geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid + geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xquat_in, geom1, geom2, worldid ) or pairid[1] >= 0 ): @@ -591,7 +592,7 @@ def sap_broadphase(m: Model, d: Data): d.nworld, d.naconmax, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, sort_index.reshape((-1, m.ngeom)), cumulative_sum.reshape(-1), nsweep, @@ -614,7 +615,7 @@ def kernel( # Data in: naconmax_in: int, geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), # Data out: collision_pair_out: wp.array(dtype=wp.vec2i), collision_pairid_out: wp.array(dtype=wp.vec2i), @@ -629,7 +630,7 @@ def kernel( if ( wp.static(_broadphase_filter(opt_broadphase_filter, ngeom_aabb, ngeom_rbound, ngeom_margin))( - geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, geom1, geom2, worldid + geom_aabb, geom_rbound, geom_margin, geom_xpos_in, geom_xquat_in, geom1, geom2, worldid ) or nxn_pairid[elementid][1] >= 0 ): @@ -676,7 +677,7 @@ def nxn_broadphase(m: Model, d: Data): m.nxn_pairid_filtered, d.naconmax, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, ], outputs=[ d.collision_pair, diff --git a/mujoco_warp/_src/collision_gjk_test.py b/mujoco_warp/_src/collision_gjk_test.py index f3635e8d2..37d4c0edf 100644 --- a/mujoco_warp/_src/collision_gjk_test.py +++ b/mujoco_warp/_src/collision_gjk_test.py @@ -29,6 +29,8 @@ from .types import MJ_MAX_EPAHORIZON from .warp_util import nested_kernel +from .math import quat_to_mat + def _geom_dist( m: Model, @@ -88,7 +90,7 @@ def _ccd_kernel( mesh_polymap: wp.array(dtype=int), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), # In: gid1: int, gid2: int, @@ -129,7 +131,7 @@ def _ccd_kernel( else: geom1.pos = pos1 if wp.static(mat1 == None): - geom1.rot = geom_xmat_in[0, gid1] + geom1.rot = quat_to_mat(geom_xquat_in[0, gid1]) else: geom1.rot = mat1 geom1.size = geom_size[0, gid1] @@ -160,7 +162,7 @@ def _ccd_kernel( else: geom2.pos = pos2 if wp.static(mat2 == None): - geom2.rot = geom_xmat_in[0, gid2] + geom2.rot = quat_to_mat(geom_xquat_in[0, gid2]) else: geom2.rot = mat2 geom2.size = geom_size[0, gid2] @@ -266,7 +268,7 @@ def _ccd_kernel( m.mesh_polymapnum, m.mesh_polymap, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, gid1, gid2, m.opt.ccd_iterations, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 428de0cb0..d8c0bb02a 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -32,6 +32,7 @@ from .math import make_frame from .math import safe_div from .math import upper_trid_index +from .math import quat_to_mat from .types import MJ_MINMU from .types import MJ_MINVAL from .types import ContactType @@ -95,7 +96,7 @@ def geom_collision_pair( mesh_polymap: wp.array(dtype=int), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), # In: geoms: wp.vec2i, worldid: int, @@ -109,12 +110,12 @@ def geom_collision_pair( geom_type2 = geom_type[g2] geom1.pos = geom_xpos_in[worldid, g1] - geom1.rot = geom_xmat_in[worldid, g1] + geom1.rot = quat_to_mat(geom_xquat_in[worldid, g1]) geom1.size = geom_size[worldid % geom_size.shape[0], g1] geom1.normal = wp.vec3(geom1.rot[0, 2], geom1.rot[1, 2], geom1.rot[2, 2]) # plane geom2.pos = geom_xpos_in[worldid, g2] - geom2.rot = geom_xmat_in[worldid, g2] + geom2.rot = quat_to_mat(geom_xquat_in[worldid, g2]) geom2.size = geom_size[worldid % geom_size.shape[0], g2] geom2.normal = wp.vec3(geom2.rot[0, 2], geom2.rot[1, 2], geom2.rot[2, 2]) # plane @@ -1607,7 +1608,7 @@ def _primitive_narrowphase( pair_friction: wp.array2d(dtype=vec5), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), naconmax_in: int, collision_pair_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array(dtype=wp.vec2i), @@ -1678,7 +1679,7 @@ def _primitive_narrowphase( mesh_polymapnum, mesh_polymap, geom_xpos_in, - geom_xmat_in, + geom_xquat_in, geoms, worldid, ) @@ -1791,7 +1792,7 @@ def primitive_narrowphase(m: Model, d: Data): m.pair_gap, m.pair_friction, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.naconmax, d.collision_pair, d.collision_pairid, diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index 027a4a3e4..b43139fd6 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -662,7 +662,7 @@ def _sdf_narrowphase( geom_plugin_index: wp.array(dtype=int), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), naconmax_in: int, collision_pair_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array(dtype=wp.vec2i), @@ -739,7 +739,7 @@ def _sdf_narrowphase( mesh_polymapnum, mesh_polymap, geom_xpos_in, - geom_xmat_in, + geom_xquat_in, geoms, worldid, ) @@ -907,7 +907,7 @@ def sdf_narrowphase(m: Model, d: Data): m.plugin_attr, m.geom_plugin_index, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.naconmax, d.collision_pair, d.collision_pairid, diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 723a29417..530dd6971 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -640,6 +640,10 @@ def make_data( mujoco.mj_kinematics(mjm, mjd) d.geom_xpos = wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3) d.geom_xmat = wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33) + geom_xquat = np.zeros((mjm.ngeom, 4)) + for i in range(mjm.ngeom): + mujoco.mju_mat2Quat(geom_xquat[i], mjd.geom_xmat[i]) + d.geom_xquat = wp.array(np.tile(geom_xquat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.quat) return d diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index 0464ffae7..1a165ffbe 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -18,6 +18,7 @@ import warp as wp from .math import safe_div +from .math import quat_to_mat from .types import MJ_MINVAL from .types import Data from .types import GeomType @@ -644,7 +645,7 @@ def _ray_geom_mesh( mat_rgba: wp.array2d(dtype=wp.vec4), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, pnt: wp.vec3, @@ -667,7 +668,7 @@ def _ray_geom_mesh( bodyexclude, ): pos = geom_xpos_in[worldid, geomid] - mat = geom_xmat_in[worldid, geomid] + mat = quat_to_mat(geom_xquat_in[worldid, geomid]) type = geom_type[geomid] if type == GeomType.MESH: @@ -729,7 +730,7 @@ def _ray( mat_rgba: wp.array2d(dtype=wp.vec4), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), # In: pnt: wp.array2d(dtype=wp.vec3), vec: wp.array2d(dtype=wp.vec3), @@ -771,7 +772,7 @@ def _ray( hfield_data, mat_rgba, geom_xpos_in, - geom_xmat_in, + geom_xquat_in, worldid, pnt[worldid, rayid], vec[worldid, rayid], @@ -875,7 +876,7 @@ def rays( m.hfield_data, m.mat_rgba, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, pnt, vec, geomgroup, From 08351177c6c54018524d37e76b51b1d03da26318 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 17:38:58 +0100 Subject: [PATCH 024/107] no more geom_xmat in mjwarp Signed-off-by: Alain Denzler --- mujoco_warp/_src/forward_test.py | 10 ++++++++-- mujoco_warp/_src/io.py | 7 +++++-- mujoco_warp/_src/io_test.py | 2 +- mujoco_warp/_src/smooth.py | 4 +--- mujoco_warp/_src/types.py | 2 -- 5 files changed, 15 insertions(+), 10 deletions(-) diff --git a/mujoco_warp/_src/forward_test.py b/mujoco_warp/_src/forward_test.py index 3d00d414e..64c48b8fe 100644 --- a/mujoco_warp/_src/forward_test.py +++ b/mujoco_warp/_src/forward_test.py @@ -319,7 +319,6 @@ def test_step1(self, xml, energy): "xanchor", "xaxis", "geom_xpos", - "geom_xmat", "site_xmat", "subtree_com", "cinert", @@ -389,7 +388,7 @@ def _getattr(arr): d_arr, is_nefc = _getattr(arr) d_arr = d_arr.numpy()[0] mjd_arr = getattr(mjd, arr) - if arr in ["geom_xmat", "site_xmat", "cam_xmat"]: + if arr in ["site_xmat", "cam_xmat"]: mjd_arr = mjd_arr.reshape(-1) d_arr = d_arr.reshape(-1) elif arr == "qM": @@ -436,6 +435,13 @@ def _getattr(arr): mujoco.mju_quat2Mat(ximat[i], xiquat[i]) _assert_eq(ximat.reshape(-1), mjd.ximat.reshape(-1), "xiquat->ximat") + # compare geom_xquat (warp) to geom_xmat (mujoco) by converting quaternion to matrix + geom_xquat = d.geom_xquat.numpy()[0] + geom_xmat = np.zeros((geom_xquat.shape[0], 9)) + for i in range(geom_xquat.shape[0]): + mujoco.mju_quat2Mat(geom_xmat[i], geom_xquat[i]) + _assert_eq(geom_xmat.reshape(-1), mjd.geom_xmat.reshape(-1), "geom_xquat->geom_xmat") + # TODO(team): sensor_pos # TODO(team): sensor_vel diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 530dd6971..f9ec5ca4d 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -617,7 +617,6 @@ def make_data( "qM": None, "qLD": None, "geom_xpos": None, - "geom_xmat": None, } for f in dataclasses.fields(types.Data): if f.name in d_kwargs: @@ -914,6 +913,11 @@ def get_data_into( mujoco.mju_quat2Mat(ximat[i], d.xiquat.numpy()[world_id, i]) result.ximat[:] = ximat.reshape((-1, 9)) + geom_xmat = np.zeros((mjm.ngeom, 9)) + for i in range(mjm.ngeom): + mujoco.mju_quat2Mat(geom_xmat[i], d.geom_xquat.numpy()[world_id, i]) + result.geom_xmat[:] = geom_xmat.reshape((-1, 9)) + result.solver_niter[0] = d.solver_niter.numpy()[world_id] result.ncon = ncon result.ne = ne @@ -953,7 +957,6 @@ def get_data_into( result.xanchor[:] = d.xanchor.numpy()[world_id] result.xaxis[:] = d.xaxis.numpy()[world_id] result.geom_xpos[:] = d.geom_xpos.numpy()[world_id] - result.geom_xmat[:] = d.geom_xmat.numpy()[world_id].reshape((-1, 9)) result.site_xpos[:] = d.site_xpos.numpy()[world_id] result.site_xmat[:] = d.site_xmat.numpy()[world_id].reshape((-1, 9)) result.cam_xpos[:] = d.cam_xpos.numpy()[world_id] diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index a6994e972..4ae278733 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -131,7 +131,7 @@ def test_get_data_into(self, nworld, world_id): "xanchor", "xaxis", "geom_xpos", - "geom_xmat", + #"geom_xmat", "site_xpos", "site_xmat", "cam_xpos", diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 4f899da34..44ba61594 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -173,7 +173,6 @@ def _geom_local_to_global( xquat_in: wp.array2d(dtype=wp.quat), # Data out: geom_xpos_out: wp.array2d(dtype=wp.vec3), - geom_xmat_out: wp.array2d(dtype=wp.mat33), geom_xquat_out: wp.array2d(dtype=wp.quat), ): worldid, geomid = wp.tid() @@ -187,7 +186,6 @@ def _geom_local_to_global( xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] geom_xpos_out[worldid, geomid] = xpos + math.rot_vec_quat(geom_pos[worldid % geom_pos.shape[0], geomid], xquat) - geom_xmat_out[worldid, geomid] = math.quat_to_mat(math.mul_quat(xquat, geom_quat[worldid % geom_quat.shape[0], geomid])) geom_xquat_out[worldid, geomid] = xquat @@ -315,7 +313,7 @@ def kinematics(m: Model, d: Data): _geom_local_to_global, dim=(d.nworld, m.ngeom), inputs=[m.body_rootid, m.body_weldid, m.body_mocapid, m.geom_bodyid, m.geom_pos, m.geom_quat, d.xpos, d.xquat], - outputs=[d.geom_xpos, d.geom_xmat, d.geom_xquat], + outputs=[d.geom_xpos, d.geom_xquat], ) wp.launch( diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 7c36559fc..2fb50b780 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1570,7 +1570,6 @@ class Data: xanchor: Cartesian position of joint anchor (nworld, njnt, 3) xaxis: Cartesian joint axis (nworld, njnt, 3) geom_xpos: Cartesian geom position (nworld, ngeom, 3) - geom_xmat: Cartesian geom orientation (nworld, ngeom, 3, 3) site_xpos: Cartesian site position (nworld, nsite, 3) site_xmat: Cartesian site orientation (nworld, nsite, 3, 3) cam_xpos: Cartesian camera position (nworld, ncam, 3) @@ -1670,7 +1669,6 @@ class Data: xanchor: array("nworld", "njnt", wp.vec3) xaxis: array("nworld", "njnt", wp.vec3) geom_xpos: array("nworld", "ngeom", wp.vec3) - geom_xmat: array("nworld", "ngeom", wp.mat33) site_xpos: array("nworld", "nsite", wp.vec3) site_xmat: array("nworld", "nsite", wp.mat33) cam_xpos: array("nworld", "ncam", wp.vec3) From 62068e165aa8b361b7219d1949f5304915fa50ff Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 17:41:16 +0100 Subject: [PATCH 025/107] maybe this helps Signed-off-by: Alain Denzler --- mujoco_warp/_src/forward_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/forward_test.py b/mujoco_warp/_src/forward_test.py index 64c48b8fe..ffdb85fe2 100644 --- a/mujoco_warp/_src/forward_test.py +++ b/mujoco_warp/_src/forward_test.py @@ -370,7 +370,7 @@ def _getattr(arr): return getattr(d, arr), False for arr in step1_field: - if arr in ("geom_xpos", "geom_xmat"): + if arr in ("geom_xpos", "geom_xmat", "geom_xquat"): # leave geom_xpos and geom_xmat untouched because they have static data continue attr, _ = _getattr(arr) From 358f4428a83a6331b2d14b1b9ed84783105ba8ae Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 17:54:35 +0100 Subject: [PATCH 026/107] more quats instead of matrices Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_convex.py | 2 +- mujoco_warp/_src/collision_driver_test.py | 2 +- mujoco_warp/_src/collision_gjk.py | 39 ++++++++--------- mujoco_warp/_src/collision_gjk_test.py | 16 +++---- mujoco_warp/_src/collision_primitive.py | 52 ++++++++++++++--------- mujoco_warp/_src/collision_sdf.py | 17 +++++--- 6 files changed, 71 insertions(+), 57 deletions(-) diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 3f7e97813..b4f0ca8c3 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -641,7 +641,7 @@ def ccd_kernel( x1_ = wp.vec3(0.0, 0.0, 0.0) for i in range(6): x1_ += prism[i] - x1 += geom1.rot @ (x1_ / 6.0) + x1 += quat_to_mat(geom1.rot) @ (x1_ / 6.0) dist, ncontact, w1, w2, idx = ccd( opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]], diff --git a/mujoco_warp/_src/collision_driver_test.py b/mujoco_warp/_src/collision_driver_test.py index 585ca0e5d..ed9251f4c 100644 --- a/mujoco_warp/_src/collision_driver_test.py +++ b/mujoco_warp/_src/collision_driver_test.py @@ -629,7 +629,7 @@ def test_plane_meshtet(self): # tetrahedron, separated in z by 0.1 convex = Geom() convex.pos = wp.vec3(0.0) - convex.rot = wp.mat33(np.eye(3)) + convex.rot = wp.quat(1.0, 0.0, 0.0, 0.0) convex.graphadr = -1 convex.vertnum = 4 convex.vertadr = 0 diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index 28fad0dfd..a9a345ceb 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -21,6 +21,7 @@ from .types import GeomType from .types import mat43 from .types import mat63 +from .math import quat_to_mat # TODO(team): improve compile time to enable backward pass wp.set_module_options({"enable_backward": False}) @@ -99,11 +100,11 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: sp.point = geom.pos + (0.5 * geom.margin) * geom.size[0] * dir return sp - local_dir = wp.transpose(geom.rot) @ dir + local_dir = wp.transpose(quat_to_mat(geom.rot)) @ dir if geomtype == GeomType.BOX: tmp = wp.sign(local_dir) res = wp.cw_mul(tmp, geom.size) - sp.point = geom.rot @ res + geom.pos + sp.point = quat_to_mat(geom.rot) @ res + geom.pos sp.vertex_index = wp.where(tmp[0] > 0, 1, 0) sp.vertex_index += wp.where(tmp[1] > 0, 2, 0) sp.vertex_index += wp.where(tmp[2] > 0, 4, 0) @@ -111,13 +112,13 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: res = local_dir * geom.size[0] # add cylinder contribution res[2] += wp.sign(local_dir[2]) * geom.size[1] - sp.point = geom.rot @ res + geom.pos + sp.point = quat_to_mat(geom.rot) @ res + geom.pos elif geomtype == GeomType.ELLIPSOID: res = wp.cw_mul(local_dir, geom.size) res = wp.normalize(res) # transform to ellipsoid res = wp.cw_mul(res, geom.size) - sp.point = geom.rot @ res + geom.pos + sp.point = quat_to_mat(geom.rot) @ res + geom.pos elif geomtype == GeomType.CYLINDER: res = wp.vec3(0.0, 0.0, 0.0) # set result in XY plane: support on circle @@ -128,7 +129,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: res[1] = local_dir[1] * scl # set result in Z direction res[2] = wp.sign(local_dir[2]) * geom.size[1] - sp.point = geom.rot @ res + geom.pos + sp.point = quat_to_mat(geom.rot) @ res + geom.pos elif geomtype == GeomType.MESH: max_dist = float(FLOAT_MIN) if geom.graphadr == -1 or geom.vertnum < 10: @@ -170,7 +171,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: sp.vertex_index = geom.graph[vert_globalid + imax] sp.point = geom.vert[geom.vertadr + sp.vertex_index] - sp.point = geom.rot @ sp.point + geom.pos + sp.point = quat_to_mat(geom.rot) @ sp.point + geom.pos elif geomtype == GeomType.HFIELD: max_dist = float(FLOAT_MIN) for i in range(6): @@ -179,7 +180,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: if dist > max_dist: max_dist = dist sp.point = vert - sp.point = geom.rot @ sp.point + geom.pos + sp.point = quat_to_mat(geom.rot) @ sp.point + geom.pos if geom.margin > 0.0: sp.point += dir * (0.5 * geom.margin) @@ -1953,12 +1954,12 @@ def multicontact( # get all possible face normals for each geom if geomtype1 == GeomType.BOX: - nnorms1 = _box_normals(nface1, feature_index1, geom1.rot, dir_neg, n1, idx1) + nnorms1 = _box_normals(nface1, feature_index1, quat_to_mat(geom1.rot), dir_neg, n1, idx1) elif geomtype1 == GeomType.MESH: nnorms1 = _mesh_normals( nface1, feature_index1, - geom1.rot, + quat_to_mat(geom1.rot), geom1.vertadr, geom1.mesh_polyadr, polynormal, @@ -1969,12 +1970,12 @@ def multicontact( idx1, ) if geomtype2 == GeomType.BOX: - nnorms2 = _box_normals(nface2, feature_index2, geom2.rot, dir, n2, idx2) + nnorms2 = _box_normals(nface2, feature_index2, quat_to_mat(geom2.rot), dir, n2, idx2) elif geomtype2 == GeomType.MESH: nnorms2 = _mesh_normals( nface2, feature_index2, - geom2.rot, + quat_to_mat(geom2.rot), geom2.vertadr, geom2.mesh_polyadr, polynormal, @@ -1995,12 +1996,12 @@ def multicontact( nnorms1 = 0 if geomtype1 == GeomType.BOX: nnorms1 = _box_edge_normals( - nface1, geom1.rot, geom1.pos, geom1.size, feature_vertex1[0], feature_vertex1[1], feature_index1[0], n1, endvert + nface1, quat_to_mat(geom1.rot), geom1.pos, geom1.size, feature_vertex1[0], feature_vertex1[1], feature_index1[0], n1, endvert ) elif geomtype1 == GeomType.MESH: nnorms1 = _mesh_edge_normals( nface1, - geom1.rot, + quat_to_mat(geom1.rot), geom1.pos, geom1.vertadr, geom1.mesh_polyadr, @@ -2027,12 +2028,12 @@ def multicontact( nnorms2 = 0 if geomtype2 == GeomType.BOX: nnorms2 = _box_edge_normals( - nface2, geom2.rot, geom2.pos, geom2.size, feature_vertex2[0], feature_vertex2[1], feature_index2[0], n2, endvert + nface2, quat_to_mat(geom2.rot), geom2.pos, geom2.size, feature_vertex2[0], feature_vertex2[1], feature_index2[0], n2, endvert ) elif geomtype2 == GeomType.MESH: nnorms2 = _mesh_edge_normals( nface2, - geom2.rot, + quat_to_mat(geom2.rot), geom2.pos, geom2.vertadr, geom2.mesh_polyadr, @@ -2066,10 +2067,10 @@ def multicontact( else: ind = wp.where(is_edge_contact_geom2, idx1[j], idx1[i]) if geomtype1 == GeomType.BOX: - nface1 = _box_face(geom1.rot, geom1.pos, geom1.size, ind, face1) + nface1 = _box_face(quat_to_mat(geom1.rot), geom1.pos, geom1.size, ind, face1) elif geomtype1 == GeomType.MESH: nface1 = _mesh_face( - geom1.rot, + quat_to_mat(geom1.rot), geom1.pos, geom1.vertadr, geom1.mesh_polyadr, @@ -2086,10 +2087,10 @@ def multicontact( nface2 = _set_edge(epa_vert2, endvert, face[0], i, face2) else: if geomtype2 == GeomType.BOX: - nface2 = _box_face(geom2.rot, geom2.pos, geom2.size, idx2[j], face2) + nface2 = _box_face(quat_to_mat(geom2.rot), geom2.pos, geom2.size, idx2[j], face2) elif geomtype2 == GeomType.MESH: nface2 = _mesh_face( - geom2.rot, + quat_to_mat(geom2.rot), geom2.pos, geom2.vertadr, geom2.mesh_polyadr, diff --git a/mujoco_warp/_src/collision_gjk_test.py b/mujoco_warp/_src/collision_gjk_test.py index 37d4c0edf..8156365a2 100644 --- a/mujoco_warp/_src/collision_gjk_test.py +++ b/mujoco_warp/_src/collision_gjk_test.py @@ -41,8 +41,8 @@ def _geom_dist( margin=0.0, pos1: wp.vec3 | None = None, pos2: wp.vec3 | None = None, - mat1: wp.mat33 | None = None, - mat2: wp.mat33 | None = None, + quat1: wp.quat | None = None, + quat2: wp.quat | None = None, ): # we run multiccd on static scenes so these need to be initialized nmaxpolygon = 10 if multiccd else 0 @@ -130,10 +130,10 @@ def _ccd_kernel( geom1.pos = geom_xpos_in[0, gid1] else: geom1.pos = pos1 - if wp.static(mat1 == None): - geom1.rot = quat_to_mat(geom_xquat_in[0, gid1]) + if wp.static(quat1 == None): + geom1.rot = geom_xquat_in[0, gid1] else: - geom1.rot = mat1 + geom1.rot = quat1 geom1.size = geom_size[0, gid1] geom1.margin = margin geom1.graphadr = -1 @@ -161,10 +161,10 @@ def _ccd_kernel( geom2.pos = geom_xpos_in[0, gid2] else: geom2.pos = pos2 - if wp.static(mat2 == None): - geom2.rot = quat_to_mat(geom_xquat_in[0, gid2]) + if wp.static(quat2 == None): + geom2.rot = geom_xquat_in[0, gid2] else: - geom2.rot = mat2 + geom2.rot = quat2 geom2.size = geom_size[0, gid2] geom2.margin = margin geom2.graphadr = -1 diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index d8c0bb02a..956031e14 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -52,7 +52,7 @@ @wp.struct class Geom: pos: wp.vec3 - rot: wp.mat33 + rot: wp.quat normal: wp.vec3 size: wp.vec3 margin: float @@ -110,14 +110,16 @@ def geom_collision_pair( geom_type2 = geom_type[g2] geom1.pos = geom_xpos_in[worldid, g1] - geom1.rot = quat_to_mat(geom_xquat_in[worldid, g1]) + geom1.rot = geom_xquat_in[worldid, g1] geom1.size = geom_size[worldid % geom_size.shape[0], g1] - geom1.normal = wp.vec3(geom1.rot[0, 2], geom1.rot[1, 2], geom1.rot[2, 2]) # plane + mat1 = quat_to_mat(geom1.rot) + geom1.normal = wp.vec3(mat1[0, 2], mat1[1, 2], mat1[2, 2]) # plane geom2.pos = geom_xpos_in[worldid, g2] - geom2.rot = quat_to_mat(geom_xquat_in[worldid, g2]) + geom2.rot = geom_xquat_in[worldid, g2] geom2.size = geom_size[worldid % geom_size.shape[0], g2] - geom2.normal = wp.vec3(geom2.rot[0, 2], geom2.rot[1, 2], geom2.rot[2, 2]) # plane + mat2 = quat_to_mat(geom2.rot) + geom2.normal = wp.vec3(mat2[0, 2], mat2[1, 2], mat2[2, 2]) # plane if geom_type1 == GeomType.MESH: dataid = geom_dataid[g1] @@ -185,8 +187,9 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup contact_count = int(0) # get points in the convex frame - plane_pos_local = wp.transpose(convex.rot) @ (plane_pos - convex.pos) - n = wp.transpose(convex.rot) @ plane_normal + mat_transposed = wp.transpose(quat_to_mat(convex.rot)) + plane_pos_local = mat_transposed @ (plane_pos - convex.pos) + n = mat_transposed @ plane_normal # Store indices in vec4 indices = wp.vec4i(-1, -1, -1, -1) @@ -381,7 +384,7 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup # Check if the index is unique (appears exactly once) if count == 1: pos = convex.vert[convex.vertadr + idx] - pos = convex.pos + convex.rot @ pos + pos = convex.pos + quat_to_mat(convex.rot) @ pos support = wp.dot(plane_pos_local - convex.vert[convex.vertadr + idx], n) dist = -support pos = pos - 0.5 * dist * plane_normal @@ -731,7 +734,8 @@ def sphere_capsule_wrapper( ): """Calculates one contact between a sphere and a capsule.""" # capsule axis - axis = wp.vec3(cap.rot[0, 2], cap.rot[1, 2], cap.rot[2, 2]) + mat = quat_to_mat(cap.rot) + axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) dist, pos, normal = sphere_capsule(sphere.pos, sphere.size[0], cap.pos, axis, cap.size[0], cap.size[1]) @@ -803,8 +807,10 @@ def capsule_capsule_wrapper( ): """Calculates contacts between two capsules.""" # capsule axes - cap1_axis = wp.vec3(cap1.rot[0, 2], cap1.rot[1, 2], cap1.rot[2, 2]) - cap2_axis = wp.vec3(cap2.rot[0, 2], cap2.rot[1, 2], cap2.rot[2, 2]) + mat1 = quat_to_mat(cap1.rot) + cap1_axis = wp.vec3(mat1[0, 2], mat1[1, 2], mat1[2, 2]) + mat2 = quat_to_mat(cap2.rot) + cap2_axis = wp.vec3(mat2[0, 2], mat2[1, 2], mat2[2, 2]) dist, pos, normal = capsule_capsule( cap1.pos, @@ -885,7 +891,8 @@ def plane_capsule_wrapper( ): """Calculates contacts between a capsule and a plane.""" # capsule axis - capsule_axis = wp.vec3(cap.rot[0, 2], cap.rot[1, 2], cap.rot[2, 2]) + mat = quat_to_mat(cap.rot) + capsule_axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) dist, pos, frame = plane_capsule( plane.normal, @@ -964,7 +971,7 @@ def plane_ellipsoid_wrapper( nacon_out: wp.array(dtype=int), ): """Calculates contacts between an ellipsoid and a plane.""" - dist, pos, normal = plane_ellipsoid(plane.normal, plane.pos, ellipsoid.pos, ellipsoid.rot, ellipsoid.size) + dist, pos, normal = plane_ellipsoid(plane.normal, plane.pos, ellipsoid.pos, quat_to_mat(ellipsoid.rot), ellipsoid.size) write_contact( naconmax_in, @@ -1033,7 +1040,7 @@ def plane_box_wrapper( nacon_out: wp.array(dtype=int), ): """Calculates contacts between a box and a plane.""" - dist, pos, normal = plane_box(plane.normal, plane.pos, box.pos, box.rot, box.size) + dist, pos, normal = plane_box(plane.normal, plane.pos, box.pos, quat_to_mat(box.rot), box.size) frame = make_frame(normal) for i in range(8): @@ -1176,7 +1183,8 @@ def sphere_cylinder_wrapper( ): """Calculates contacts between a sphere and a cylinder.""" # cylinder axis - cylinder_axis = wp.vec3(cylinder.rot[0, 2], cylinder.rot[1, 2], cylinder.rot[2, 2]) + mat = quat_to_mat(cylinder.rot) + cylinder_axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) dist, pos, normal = sphere_cylinder( sphere.pos, @@ -1255,7 +1263,8 @@ def plane_cylinder_wrapper( ): """Calculates contacts between a cylinder and a plane.""" # cylinder axis - cylinder_axis = wp.vec3(cylinder.rot[0, 2], cylinder.rot[1, 2], cylinder.rot[2, 2]) + mat = quat_to_mat(cylinder.rot) + cylinder_axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) dist, pos, normal = plane_cylinder( plane.normal, @@ -1334,7 +1343,7 @@ def sphere_box_wrapper( contact_geomcollisionid_out: wp.array(dtype=int), nacon_out: wp.array(dtype=int), ): - dist, pos, normal = sphere_box(sphere.pos, sphere.size[0], box.pos, box.rot, box.size) + dist, pos, normal = sphere_box(sphere.pos, sphere.size[0], box.pos, quat_to_mat(box.rot), box.size) write_contact( naconmax_in, @@ -1404,7 +1413,8 @@ def capsule_box_wrapper( ): """Calculates contacts between a capsule and a box.""" # Extract capsule axis - axis = wp.vec3(cap.rot[0, 2], cap.rot[1, 2], cap.rot[2, 2]) + mat = quat_to_mat(cap.rot) + axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) # Call the core function to get contact geometry dist, pos, normal = capsule_box( @@ -1413,7 +1423,7 @@ def capsule_box_wrapper( cap.size[0], # capsule radius cap.size[1], # capsule half length box.pos, - box.rot, + quat_to_mat(box.rot), box.size, ) @@ -1489,10 +1499,10 @@ def box_box_wrapper( # Call the core function to get contact geometry dist, pos, normal = box_box( box1.pos, - box1.rot, + quat_to_mat(box1.rot), box1.size, box2.pos, - box2.rot, + quat_to_mat(box2.rot), box2.size, margin, ) diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index b43139fd6..bb6617f06 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -21,6 +21,7 @@ from .collision_primitive import geom_collision_pair from .collision_primitive import write_contact from .math import make_frame +from .math import quat_to_mat from .ray import ray_mesh from .types import Data from .types import GeomType @@ -751,8 +752,10 @@ def _sdf_narrowphase( g1_plugin = geom_plugin_index[g1] g2_plugin = geom_plugin_index[g2] - g1_to_g2_rot = wp.transpose(geom1.rot) * geom2.rot - g1_to_g2_pos = wp.transpose(geom1.rot) * (geom2.pos - geom1.pos) + mat_rot_1 = quat_to_mat(geom1.rot) + mat_rot_2 = quat_to_mat(geom2.rot) + g1_to_g2_rot = wp.transpose(mat_rot_1) * mat_rot_2 + g1_to_g2_pos = wp.transpose(mat_rot_1) * (geom2.pos - geom1.pos) aabb_pos = geom_aabb[aabb_id, g1, 0] aabb_size = geom_aabb[aabb_id, g1, 1] identity = wp.identity(3, dtype=float) @@ -765,9 +768,9 @@ def _sdf_narrowphase( aabb_intersection.max = wp.min(aabb1.max, aabb2.max) pos2 = geom2.pos - rot2 = geom2.rot + rot2 = mat_rot_2 pos1 = geom1.pos - rot1 = geom1.rot + rot1 = mat_rot_1 attr1, g1_plugin_id, volume_data1, mesh_data1 = get_sdf_params( oct_child, oct_aabb, oct_coeff, plugin, plugin_attr, type1, geom1.size, g1_plugin, geom_dataid[g1] @@ -784,7 +787,7 @@ def _sdf_narrowphase( mesh_data1.mesh_face = mesh_face mesh_data1.data_id = geom_dataid[g1] mesh_data1.pos = geom1.pos - mesh_data1.mat = geom1.rot + mesh_data1.mat = mat_rot_1 mesh_data1.pnt = wp.vec3(-1.0) mesh_data1.vec = wp.vec3(0.0) mesh_data1.valid = True @@ -796,7 +799,7 @@ def _sdf_narrowphase( mesh_data2.mesh_face = mesh_face mesh_data2.data_id = geom_dataid[g2] mesh_data2.pos = geom2.pos - mesh_data2.mat = geom2.rot + mesh_data2.mat = mat_rot_2 mesh_data2.pnt = wp.vec3(-1.0) mesh_data2.vec = wp.vec3(0.0) mesh_data2.valid = True @@ -806,7 +809,7 @@ def _sdf_narrowphase( aabb_intersection.min[1] + (aabb_intersection.max[1] - aabb_intersection.min[1]) * halton(i, 3), aabb_intersection.min[2] + (aabb_intersection.max[2] - aabb_intersection.min[2]) * halton(i, 5), ) - x = geom1.rot * x_g2 + geom1.pos + x = mat_rot_1 * x_g2 + geom1.pos x0_initial = wp.transpose(rot2) * (x - pos2) dist, pos, n = gradient_descent( type1, From 7ea3e42e66a7d07084e999f60adde7015b658a02 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 17:57:58 +0100 Subject: [PATCH 027/107] remove unnecessary line Signed-off-by: Alain Denzler --- mujoco_warp/_src/io.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index f9ec5ca4d..8310c8ee4 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -638,7 +638,6 @@ def make_data( mjd = mujoco.MjData(mjm) mujoco.mj_kinematics(mjm, mjd) d.geom_xpos = wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3) - d.geom_xmat = wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33) geom_xquat = np.zeros((mjm.ngeom, 4)) for i in range(mjm.ngeom): mujoco.mju_mat2Quat(geom_xquat[i], mjd.geom_xmat[i]) From 3a2469febb75fd35bd00e8f33c6a6c679073f21d Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 8 Dec 2025 18:12:31 +0100 Subject: [PATCH 028/107] fix the GJK test Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk_test.py | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk_test.py b/mujoco_warp/_src/collision_gjk_test.py index 8156365a2..45584139a 100644 --- a/mujoco_warp/_src/collision_gjk_test.py +++ b/mujoco_warp/_src/collision_gjk_test.py @@ -13,6 +13,7 @@ # limitations under the License. # ============================================================================== +import mujoco import numpy as np import warp as wp from absl.testing import absltest @@ -558,7 +559,7 @@ def test_cylinder_box(self): ) pos = wp.vec3(0.00015228791744448245, -0.00074981129728257656, 0.29839199781417846680) - rot = wp.mat33( + rot_mat = np.array([ 0.99996972084045410156, 0.00776371126994490623, -0.00043433305108919740, @@ -568,9 +569,12 @@ def test_cylinder_box(self): 0.00043175052269361913, 0.00033431366318836808, 0.99999988079071044922, - ) + ]) + rot_quat = np.zeros(4) + mujoco.mju_mat2Quat(rot_quat, rot_mat) + rot = wp.quat(rot_quat[0], rot_quat[1], rot_quat[2], rot_quat[3]) - dist, _, _, _ = _geom_dist(m, d, 0, 1, pos2=pos, mat2=rot) + dist, _, _, _ = _geom_dist(m, d, 0, 1, pos2=pos, quat2=rot) self.assertAlmostEqual(dist, -0.0016624178339902445) def test_box_box_float(self): @@ -587,7 +591,7 @@ def test_box_box_float(self): ) pos1 = wp.vec3(-0.17624500393867492676, -0.12375499308109283447, 0.12499777972698211670) - rot1 = wp.mat33( + rot1_mat = np.array([ 1.00000000000000000000, -0.00000000184385418045, -0.00000025833372774287, @@ -597,10 +601,13 @@ def test_box_box_float(self): 0.00000025833372774287, -0.00000024928382913458, 1.0000000000000000000, - ) + ]) + rot1_quat = np.zeros(4) + mujoco.mju_mat2Quat(rot1_quat, rot1_mat) + rot1 = wp.quat(rot1_quat[0], rot1_quat[1], rot1_quat[2], rot1_quat[3]) pos2 = wp.vec3(-0.17624500393867492676, -0.12375499308109283447, 0.17499557137489318848) - rot2 = wp.mat33( + rot2_mat = np.array([ 1.00000000000000000000, -0.00000000184292525685, 0.00000012980596864054, @@ -610,9 +617,12 @@ def test_box_box_float(self): -0.00000012980596864054, 0.00000014602545661546, 1.00000000000000000000, - ) + ]) + rot2_quat = np.zeros(4) + mujoco.mju_mat2Quat(rot2_quat, rot2_mat) + rot2 = wp.quat(rot2_quat[0], rot2_quat[1], rot2_quat[2], rot2_quat[3]) - dist, ncon, _, _ = _geom_dist(m, d, 0, 1, multiccd=False, pos1=pos1, mat1=rot1, pos2=pos2, mat2=rot2) + dist, ncon, _, _ = _geom_dist(m, d, 0, 1, multiccd=False, pos1=pos1, quat1=rot1, pos2=pos2, quat2=rot2) self.assertEqual(ncon, 1) self.assertLess(dist, 0.0001) # real depth is ~ 2E-6 From 4a89cbb843f3c154fc2d29562eb7b994b73ed3ee Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 11:57:14 +0100 Subject: [PATCH 029/107] fix Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 44ba61594..29b0978a6 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -186,7 +186,7 @@ def _geom_local_to_global( xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] geom_xpos_out[worldid, geomid] = xpos + math.rot_vec_quat(geom_pos[worldid % geom_pos.shape[0], geomid], xquat) - geom_xquat_out[worldid, geomid] = xquat + geom_xquat_out[worldid, geomid] = math.mul_quat(xquat, geom_quat[worldid % geom_quat.shape[0], geomid]) @wp.kernel From 6c2d5e2aeab7d744c38747b69a394ccfc39ea4af Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:11:39 +0100 Subject: [PATCH 030/107] add site_xquat Signed-off-by: Alain Denzler --- mujoco_warp/_src/types.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 28842f024..15c579194 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1641,6 +1641,7 @@ class Data: ncollision: collision count from broadphase (1,) xiquat: Cartesian orientation of body inertia (nworld, nbody, 4) geom_xquat: Cartesian orientation of geom (nworld, ngeom, 4) + site_xquat: Cartesian orientation of site (nworld, nsite, 4) """ solver_niter: array("nworld", int) @@ -1739,4 +1740,5 @@ class Data: xiquat: array("nworld", "nbody", wp.quat) geom_xquat: array("nworld", "ngeom", wp.quat) + site_xquat: array("nworld", "nsite", wp.quat) From f7692a51de38560e88c25a239527afe17e0deefc Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:15:35 +0100 Subject: [PATCH 031/107] io.py Signed-off-by: Alain Denzler --- mujoco_warp/_src/io.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index f46a71157..39e4108a7 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -788,6 +788,7 @@ def put_data( "nsolving": None, "xiquat": None, "geom_xquat": None, + "site_xquat": None, } for f in dataclasses.fields(types.Data): if f.name in d_kwargs: @@ -834,7 +835,7 @@ def put_data( d.actuator_moment = wp.array(np.full((nworld, mjm.nu, mjm.nv), actuator_moment), dtype=float) # convert xmat fields to xquat - for xmat, attr, n in [(mjd.ximat, "xiquat", mjm.nbody), (mjd.geom_xmat, "geom_xquat", mjm.ngeom)]: + for xmat, attr, n in [(mjd.ximat, "xiquat", mjm.nbody), (mjd.geom_xmat, "geom_xquat", mjm.ngeom), (mjd.site_xmat, "site_xquat", mjm.nsite)]: xquat = np.zeros((n, 4)) for i in range(n): mujoco.mju_mat2Quat(xquat[i], xmat[i]) @@ -913,12 +914,17 @@ def get_data_into( ximat = np.zeros((mjm.nbody, 9)) for i in range(mjm.nbody): mujoco.mju_quat2Mat(ximat[i], d.xiquat.numpy()[world_id, i]) - result.ximat[:] = ximat.reshape((-1, 9)) + result.ximat[:] = ximat geom_xmat = np.zeros((mjm.ngeom, 9)) for i in range(mjm.ngeom): mujoco.mju_quat2Mat(geom_xmat[i], d.geom_xquat.numpy()[world_id, i]) - result.geom_xmat[:] = geom_xmat.reshape((-1, 9)) + result.geom_xmat[:] = geom_xmat + + site_xmat = np.zeros((mjm.nsite, 9)) + for i in range(mjm.nsite): + mujoco.mju_quat2Mat(site_xmat[i], d.site_xquat.numpy()[world_id, i]) + result.site_xmat[:] = site_xmat result.solver_niter[0] = d.solver_niter.numpy()[world_id] result.ncon = ncon From 27984f9e697f921286b8e9f3edc4e6df17fe6c6f Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:18:19 +0100 Subject: [PATCH 032/107] add initialization Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 29b0978a6..d3f5743b2 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -201,6 +201,7 @@ def _site_local_to_global( # Data out: site_xpos_out: wp.array2d(dtype=wp.vec3), site_xmat_out: wp.array2d(dtype=wp.mat33), + site_xquat_out: wp.array2d(dtype=wp.quat), ): worldid, siteid = wp.tid() bodyid = site_bodyid[siteid] @@ -208,7 +209,7 @@ def _site_local_to_global( xquat = xquat_in[worldid, bodyid] site_xpos_out[worldid, siteid] = xpos + math.rot_vec_quat(site_pos[worldid % site_pos.shape[0], siteid], xquat) site_xmat_out[worldid, siteid] = math.quat_to_mat(math.mul_quat(xquat, site_quat[worldid % site_quat.shape[0], siteid])) - + site_xquat_out[worldid, siteid] = math.mul_quat(xquat, site_quat[worldid % site_quat.shape[0], siteid]) @wp.kernel def _flex_vertices( @@ -320,7 +321,7 @@ def kinematics(m: Model, d: Data): _site_local_to_global, dim=(d.nworld, m.nsite), inputs=[m.site_bodyid, m.site_pos, m.site_quat, d.xpos, d.xquat], - outputs=[d.site_xpos, d.site_xmat], + outputs=[d.site_xpos, d.site_xmat, d.site_xquat], ) From fb417cf8f98b7368b4fbd25a11d83370ef2f4352 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:18:54 +0100 Subject: [PATCH 033/107] change the forward test Signed-off-by: Alain Denzler --- mujoco_warp/_src/forward_test.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/forward_test.py b/mujoco_warp/_src/forward_test.py index dce0f43aa..b29959d94 100644 --- a/mujoco_warp/_src/forward_test.py +++ b/mujoco_warp/_src/forward_test.py @@ -319,7 +319,6 @@ def test_step1(self, xml, energy): "xanchor", "xaxis", "geom_xpos", - "site_xmat", "subtree_com", "cinert", "cdof", @@ -388,7 +387,7 @@ def _getattr(arr): d_arr, is_nefc = _getattr(arr) d_arr = d_arr.numpy()[0] mjd_arr = getattr(mjd, arr) - if arr in ["site_xmat", "cam_xmat"]: + if arr in ["cam_xmat"]: mjd_arr = mjd_arr.reshape(-1) d_arr = d_arr.reshape(-1) elif arr == "qM": @@ -442,6 +441,13 @@ def _getattr(arr): mujoco.mju_quat2Mat(geom_xmat[i], geom_xquat[i]) _assert_eq(geom_xmat.reshape(-1), mjd.geom_xmat.reshape(-1), "geom_xquat->geom_xmat") + # compare site_xquat (warp) to site_xmat (mujoco) by converting quaternion to matrix + site_xquat = d.site_xquat.numpy()[0] + site_xmat = np.zeros((site_xquat.shape[0], 9)) + for i in range(site_xquat.shape[0]): + mujoco.mju_quat2Mat(site_xmat[i], site_xquat[i]) + _assert_eq(site_xmat.reshape(-1), mjd.site_xmat.reshape(-1), "site_xquat->site_xmat") + # TODO(team): sensor_pos # TODO(team): sensor_vel From 7517cac39c0b59b37dafa4fb799f3bb331405d9f Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:19:36 +0100 Subject: [PATCH 034/107] fix get_data_into test Signed-off-by: Alain Denzler --- mujoco_warp/_src/io_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index 4ae278733..d80c66edb 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -133,7 +133,7 @@ def test_get_data_into(self, nworld, world_id): "geom_xpos", #"geom_xmat", "site_xpos", - "site_xmat", + #"site_xmat", "cam_xpos", "cam_xmat", "light_xpos", From 1c03048c4b3bdebc6d7a63f412a0c708e29fea38 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:19:51 +0100 Subject: [PATCH 035/107] fix in io.py Signed-off-by: Alain Denzler --- mujoco_warp/_src/io.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 39e4108a7..f78c80692 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -966,7 +966,6 @@ def get_data_into( result.xaxis[:] = d.xaxis.numpy()[world_id] result.geom_xpos[:] = d.geom_xpos.numpy()[world_id] result.site_xpos[:] = d.site_xpos.numpy()[world_id] - result.site_xmat[:] = d.site_xmat.numpy()[world_id].reshape((-1, 9)) result.cam_xpos[:] = d.cam_xpos.numpy()[world_id] result.cam_xmat[:] = d.cam_xmat.numpy()[world_id].reshape((-1, 9)) result.light_xpos[:] = d.light_xpos.numpy()[world_id] From 86fffa735a51cee22e3a199a6d91900cc14dab23 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:29:02 +0100 Subject: [PATCH 036/107] transmission Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index d3f5743b2..b7769cde5 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -1778,7 +1778,7 @@ def _transmission( qpos_in: wp.array2d(dtype=float), xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cdof_in: wp.array2d(dtype=wp.spatial_vector), ten_J_in: wp.array3d(dtype=float), @@ -1835,7 +1835,7 @@ def _transmission( idslider = trnid[1] gear0 = gear[0] rod = actuator_cranklength[actid] - site_xmat = site_xmat_in[worldid, idslider] + site_xmat = math.quat_to_mat(site_xquat_in[worldid, idslider]) axis = wp.vec3(site_xmat[0, 2], site_xmat[1, 2], site_xmat[2, 2]) site_xpos_id = site_xpos_in[worldid, id] site_xpos_idslider = site_xpos_in[worldid, idslider] @@ -1923,7 +1923,7 @@ def _transmission( # reference site undefined if refid == -1: # wrench: gear expressed in global frame - site_xmat = site_xmat_in[worldid, siteid] + site_xmat = math.quat_to_mat(site_xquat_in[worldid, siteid]) wrench_translation = site_xmat @ gear_translation wrench_rotation = site_xmat @ gear_rotational @@ -1976,7 +1976,7 @@ def _transmission( site_xpos = site_xpos_in[worldid, siteid] ref_xpos = site_xpos_in[worldid, refid] - ref_xmat = site_xmat_in[worldid, refid] + ref_xmat = math.quat_to_mat(site_xquat_in[worldid, refid]) length = float(0.0) @@ -2181,7 +2181,7 @@ def transmission(m: Model, d: Data): d.qpos, d.xquat, d.site_xpos, - d.site_xmat, + d.site_xquat, d.subtree_com, d.cdof, d.ten_J, From 330d760f5d2595f8fa1e7f592c7fffe63c0c0151 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:40:58 +0100 Subject: [PATCH 037/107] sensor_pos Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 38 +++++++++++++++++------------------ mujoco_warp/_src/util_misc.py | 3 ++- 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index b79725a89..5c8e8e6ce 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -112,13 +112,13 @@ def _magnetometer( # Model: opt_magnetic: wp.array(dtype=wp.vec3), # Data in: - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, objid: int, ) -> wp.vec3: magnetic = opt_magnetic[worldid % opt_magnetic.shape[0]] - return wp.transpose(site_xmat_in[worldid, objid]) @ magnetic + return wp.transpose(math.quat_to_mat(site_xquat_in[worldid, objid])) @ magnetic @wp.func @@ -195,7 +195,7 @@ def _sensor_rangefinder_init( sensor_rangefinder_adr: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), # Out: pnt_out: wp.array2d(dtype=wp.vec3), vec_out: wp.array2d(dtype=wp.vec3), @@ -204,7 +204,7 @@ def _sensor_rangefinder_init( sensorid = sensor_rangefinder_adr[rfid] objid = sensor_objid[sensorid] site_xpos = site_xpos_in[worldid, objid] - site_xmat = site_xmat_in[worldid, objid] + site_xmat = math.quat_to_mat(site_xquat_in[worldid, objid]) pnt_out[worldid, rfid] = site_xpos vec_out[worldid, rfid] = wp.vec3(site_xmat[0, 2], site_xmat[1, 2], site_xmat[2, 2]) @@ -285,7 +285,7 @@ def _frame_pos( geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), # In: @@ -322,7 +322,7 @@ def _frame_pos( xmat_ref = math.quat_to_mat(geom_xquat_in[worldid, refid]) elif reftype == ObjType.SITE: xpos_ref = site_xpos_in[worldid, refid] - xmat_ref = site_xmat_in[worldid, refid] + xmat_ref = math.quat_to_mat(site_xquat_in[worldid, refid]) elif reftype == ObjType.CAMERA: xpos_ref = cam_xpos_in[worldid, refid] xmat_ref = cam_xmat_in[worldid, refid] @@ -340,7 +340,7 @@ def _frame_axis( xquat_in: wp.array2d(dtype=wp.quat), xiquat_in: wp.array2d(dtype=wp.quat), geom_xquat_in: wp.array2d(dtype=wp.quat), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), cam_xmat_in: wp.array2d(dtype=wp.mat33), # In: worldid: int, @@ -360,7 +360,7 @@ def _frame_axis( xmat = math.quat_to_mat(geom_xquat_in[worldid, objid]) axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) elif objtype == ObjType.SITE: - xmat = site_xmat_in[worldid, objid] + xmat = math.quat_to_mat(site_xquat_in[worldid, objid]) axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) elif objtype == ObjType.CAMERA: xmat = cam_xmat_in[worldid, objid] @@ -378,7 +378,7 @@ def _frame_axis( elif reftype == ObjType.GEOM: xmat_ref = math.quat_to_mat(geom_xquat_in[worldid, refid]) elif reftype == ObjType.SITE: - xmat_ref = site_xmat_in[worldid, refid] + xmat_ref = math.quat_to_mat(site_xquat_in[worldid, refid]) elif reftype == ObjType.CAMERA: xmat_ref = cam_xmat_in[worldid, refid] else: # UNKNOWN @@ -496,7 +496,7 @@ def _sensor_pos( geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), subtree_com_in: wp.array2d(dtype=wp.vec3), @@ -515,7 +515,7 @@ def _sensor_pos( out = sensordata_out[worldid] if sensortype == SensorType.MAGNETOMETER: - vec3 = _magnetometer(opt_magnetic, site_xmat_in, worldid, objid) + vec3 = _magnetometer(opt_magnetic, site_xquat_in, worldid, objid) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.CAMPROJECTION: refid = sensor_refid[sensorid] @@ -550,7 +550,7 @@ def _sensor_pos( geom_xpos_in, geom_xquat_in, site_xpos_in, - site_xmat_in, + site_xquat_in, cam_xpos_in, cam_xmat_in, worldid, @@ -571,7 +571,7 @@ def _sensor_pos( elif sensortype == SensorType.FRAMEZAXIS: axis = 2 vec3 = _frame_axis( - xquat_in, xiquat_in, geom_xquat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis + xquat_in, xiquat_in, geom_xquat_in, site_xquat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEQUAT: @@ -689,7 +689,7 @@ def _sensor_pos( else: return # should not occur refid = sensor_refid[sensorid] - val_bool = inside_geom(site_xpos_in[worldid, refid], site_xmat_in[worldid, refid], site_size[refid], site_type[refid], xpos) + val_bool = inside_geom(site_xpos_in[worldid, refid], site_xquat_in[worldid, refid], site_size[refid], site_type[refid], xpos) _write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, float(val_bool), out) elif sensortype == SensorType.E_POTENTIAL: val = energy_in[worldid][0] @@ -770,7 +770,7 @@ def sensor_pos(m: Model, d: Data): wp.launch( _sensor_rangefinder_init, dim=(d.nworld, m.sensor_rangefinder_adr.size), - inputs=[m.sensor_objid, m.sensor_rangefinder_adr, d.site_xpos, d.site_xmat], + inputs=[m.sensor_objid, m.sensor_rangefinder_adr, d.site_xpos, d.site_xquat], outputs=[rangefinder_pnt, rangefinder_vec], ) @@ -859,7 +859,7 @@ def sensor_pos(m: Model, d: Data): d.geom_xpos, d.geom_xquat, d.site_xpos, - d.site_xmat, + d.site_xquat, d.cam_xpos, d.cam_xmat, d.subtree_com, @@ -2236,7 +2236,7 @@ def _contact_match( sensor_contact_adr: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), contact_dist_in: wp.array(dtype=float), contact_pos_in: wp.array(dtype=wp.vec3), contact_frame_in: wp.array(dtype=wp.mat33), @@ -2276,7 +2276,7 @@ def _contact_match( # site filter if objtype == ObjType.SITE: if not inside_geom( - site_xpos_in[worldid, objid], site_xmat_in[worldid, objid], site_size[objid], site_type[objid], contact_pos_in[contactid] + site_xpos_in[worldid, objid], site_xquat_in[worldid, objid], site_size[objid], site_type[objid], contact_pos_in[contactid] ): return @@ -2491,7 +2491,7 @@ def sensor_acc(m: Model, d: Data): m.sensor_intprm, m.sensor_contact_adr, d.site_xpos, - d.site_xmat, + d.site_xquat, d.contact.dist, d.contact.pos, d.contact.frame, diff --git a/mujoco_warp/_src/util_misc.py b/mujoco_warp/_src/util_misc.py index 8c48fbe0d..8968795e5 100644 --- a/mujoco_warp/_src/util_misc.py +++ b/mujoco_warp/_src/util_misc.py @@ -599,7 +599,7 @@ def muscle_dynamics(control: float, activation: float, prm: vec10) -> float: @wp.func -def inside_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, geomtype: int, point: wp.vec3) -> bool: +def inside_geom(pos: wp.vec3, quat: wp.quat, size: wp.vec3, geomtype: int, point: wp.vec3) -> bool: """Return True if point is inside primitive geom, False otherwise.""" # vector from geom to point vec = point - pos @@ -609,6 +609,7 @@ def inside_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, geomtype: int, point return wp.dot(vec, vec) < size[0] * size[0] # rotate into local frame + mat = math.quat_to_mat(quat) plocal = wp.transpose(mat) @ vec # handle other geom types From 7a055b744f86d73918c259965fd02b65b199294c Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:43:11 +0100 Subject: [PATCH 038/107] sensor_vel Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 5c8e8e6ce..c92b86c49 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -903,7 +903,7 @@ def _velocimeter( site_bodyid: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), # In: @@ -912,7 +912,7 @@ def _velocimeter( ) -> wp.vec3: bodyid = site_bodyid[objid] pos = site_xpos_in[worldid, objid] - rot = site_xmat_in[worldid, objid] + rot = math.quat_to_mat(site_xquat_in[worldid, objid]) cvel = cvel_in[worldid, bodyid] ang = wp.spatial_top(cvel) lin = wp.spatial_bottom(cvel) @@ -926,14 +926,14 @@ def _gyro( # Model: site_bodyid: wp.array(dtype=int), # Data in: - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), cvel_in: wp.array2d(dtype=wp.spatial_vector), # In: worldid: int, objid: int, ) -> wp.vec3: bodyid = site_bodyid[objid] - rot = site_xmat_in[worldid, objid] + rot = math.quat_to_mat(site_xquat_in[worldid, objid]) cvel = cvel_in[worldid, bodyid] ang = wp.spatial_top(cvel) return wp.transpose(rot) @ ang @@ -1055,7 +1055,7 @@ def _frame_linvel( geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), subtree_com_in: wp.array2d(dtype=wp.vec3), @@ -1091,7 +1091,7 @@ def _frame_linvel( xmatref = math.quat_to_mat(geom_xquat_in[worldid, refid]) elif reftype == ObjType.SITE: xposref = site_xpos_in[worldid, refid] - xmatref = site_xmat_in[worldid, refid] + xmatref = math.quat_to_mat(site_xquat_in[worldid, refid]) elif reftype == ObjType.CAMERA: xposref = cam_xpos_in[worldid, refid] xmatref = cam_xmat_in[worldid, refid] @@ -1161,7 +1161,7 @@ def _frame_angvel( geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), subtree_com_in: wp.array2d(dtype=wp.vec3), @@ -1199,7 +1199,7 @@ def _frame_angvel( elif reftype == ObjType.GEOM: xmatref = math.quat_to_mat(geom_xquat_in[worldid, refid]) elif reftype == ObjType.SITE: - xmatref = site_xmat_in[worldid, refid] + xmatref = math.quat_to_mat(site_xquat_in[worldid, refid]) elif reftype == ObjType.CAMERA: xmatref = cam_xmat_in[worldid, refid] else: # UNKNOWN @@ -1264,7 +1264,7 @@ def _sensor_vel( geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), subtree_com_in: wp.array2d(dtype=wp.vec3), @@ -1283,10 +1283,10 @@ def _sensor_vel( out = sensordata_out[worldid] if sensortype == SensorType.VELOCIMETER: - vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, worldid, objid) + vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, site_xquat_in, subtree_com_in, cvel_in, worldid, objid) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.GYRO: - vec3 = _gyro(site_bodyid, site_xmat_in, cvel_in, worldid, objid) + vec3 = _gyro(site_bodyid, site_xquat_in, cvel_in, worldid, objid) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.JOINTVEL: val = _joint_vel(jnt_dofadr, qvel_in, worldid, objid) @@ -1316,7 +1316,7 @@ def _sensor_vel( geom_xpos_in, geom_xquat_in, site_xpos_in, - site_xmat_in, + site_xquat_in, cam_xpos_in, cam_xmat_in, subtree_com_in, @@ -1344,7 +1344,7 @@ def _sensor_vel( geom_xpos_in, geom_xquat_in, site_xpos_in, - site_xmat_in, + site_xquat_in, cam_xpos_in, cam_xmat_in, subtree_com_in, @@ -1399,7 +1399,7 @@ def sensor_vel(m: Model, d: Data): d.geom_xpos, d.geom_xquat, d.site_xpos, - d.site_xmat, + d.site_xquat, d.cam_xpos, d.cam_xmat, d.subtree_com, From d38aff10c4aa9f2ea830a2336da5b6c86f8ac2dd Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 13:55:58 +0100 Subject: [PATCH 039/107] sensor_acc Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_sdf.py | 12 +++---- mujoco_warp/_src/ray.py | 57 ++++++++++++++++--------------- mujoco_warp/_src/sensor.py | 28 +++++++-------- 3 files changed, 49 insertions(+), 48 deletions(-) diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index bb6617f06..2c65f32fe 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -69,7 +69,7 @@ class MeshData: data_id: int data_id: int pos: wp.vec3 - mat: wp.mat33 + quat: wp.quat pnt: wp.vec3 vec: wp.vec3 valid: bool = False @@ -376,7 +376,7 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume mesh_data.mesh_face, mesh_data.data_id, mesh_data.pos, - mesh_data.mat, + mesh_data.quat, mesh_data.pnt, mesh_data.vec, ) @@ -389,7 +389,7 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume mesh_data.mesh_face, mesh_data.data_id, mesh_data.pos, - mesh_data.mat, + mesh_data.quat, mesh_data.pnt, -mesh_data.vec, ) @@ -425,7 +425,7 @@ def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: V mesh_data.mesh_face, mesh_data.data_id, mesh_data.pos, - mesh_data.mat, + mesh_data.quat, mesh_data.pnt, mesh_data.vec, ) @@ -787,7 +787,7 @@ def _sdf_narrowphase( mesh_data1.mesh_face = mesh_face mesh_data1.data_id = geom_dataid[g1] mesh_data1.pos = geom1.pos - mesh_data1.mat = mat_rot_1 + mesh_data1.quat = geom1.rot mesh_data1.pnt = wp.vec3(-1.0) mesh_data1.vec = wp.vec3(0.0) mesh_data1.valid = True @@ -799,7 +799,7 @@ def _sdf_narrowphase( mesh_data2.mesh_face = mesh_face mesh_data2.data_id = geom_dataid[g2] mesh_data2.pos = geom2.pos - mesh_data2.mat = mat_rot_2 + mesh_data2.quat = geom2.rot mesh_data2.pnt = wp.vec3(-1.0) mesh_data2.vec = wp.vec3(0.0) mesh_data2.valid = True diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index 1a165ffbe..c2f30335e 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -29,7 +29,7 @@ @wp.func -def _ray_map(pos: wp.vec3, mat: wp.mat33, pnt: wp.vec3, vec: wp.vec3) -> Tuple[wp.vec3, wp.vec3]: +def _ray_map(pos: wp.vec3, quat: wp.quat, pnt: wp.vec3, vec: wp.vec3) -> Tuple[wp.vec3, wp.vec3]: """Maps ray to local geom frame coordinates. Args: @@ -41,7 +41,7 @@ def _ray_map(pos: wp.vec3, mat: wp.mat33, pnt: wp.vec3, vec: wp.vec3) -> Tuple[w Returns: 3D point and 3D direction in local geom frame """ - matT = wp.transpose(mat) + matT = wp.transpose(quat_to_mat(quat)) lpnt = matT @ (pnt - pos) lvec = matT @ vec @@ -182,10 +182,10 @@ def _ray_triangle(v0: wp.vec3, v1: wp.vec3, v2: wp.vec3, pnt: wp.vec3, vec: wp.v @wp.func -def _ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> float: +def _ray_plane(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> float: """Returns the distance at which a ray intersects with a plane.""" # map to local frame - lpnt, lvec = _ray_map(pos, mat, pnt, vec) + lpnt, lvec = _ray_map(pos, quat, pnt, vec) # z-vec not pointing towards front face: reject if lvec[2] > -MJ_MINVAL: @@ -219,7 +219,7 @@ def _ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> fl @wp.func -def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> float: +def _ray_capsule(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> float: """Returns the distance at which a ray intersects with a capsule.""" # bounding sphere test ssz = size[0] + size[1] @@ -227,7 +227,7 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: return wp.inf # map to local frame - lpnt, lvec = _ray_map(pos, mat, pnt, vec) + lpnt, lvec = _ray_map(pos, quat, pnt, vec) # init solution x = -1.0 @@ -275,10 +275,10 @@ def _ray_capsule(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: @wp.func -def _ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> float: +def _ray_ellipsoid(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> float: """Returns the distance at which a ray intersects with an ellipsoid.""" # map to local frame - lpnt, lvec = _ray_map(pos, mat, pnt, vec) + lpnt, lvec = _ray_map(pos, quat, pnt, vec) # invert size^2 s = wp.vec3(safe_div(1.0, size[0] * size[0]), safe_div(1.0, size[1] * size[1]), safe_div(1.0, size[2] * size[2])) @@ -295,7 +295,7 @@ def _ray_ellipsoid(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec @wp.func -def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> float: +def _ray_cylinder(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> float: """Returns the distance at which a ray intersects with a cylinder.""" # bounding sphere test ssz = size[0] * size[0] + size[1] * size[1] @@ -303,7 +303,7 @@ def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: return wp.inf # map to local frame - lpnt, lvec = _ray_map(pos, mat, pnt, vec) + lpnt, lvec = _ray_map(pos, quat, pnt, vec) # init solution x = wp.inf @@ -344,7 +344,7 @@ def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: @wp.func -def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6]: +def _ray_box(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6]: """Returns the distance at which a ray intersects with a box.""" all = vec6(-1.0, -1.0, -1.0, -1.0, -1.0, -1.0) @@ -354,7 +354,7 @@ def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.v return wp.inf, all # map to local frame - lpnt, lvec = _ray_map(pos, mat, pnt, vec) + lpnt, lvec = _ray_map(pos, quat, pnt, vec) # init solution x = wp.inf @@ -399,7 +399,7 @@ def _ray_hfield( hfield_data: wp.array(dtype=float), # In: pos: wp.vec3, - mat: wp.mat33, + quat: wp.quat, pnt: wp.vec3, vec: wp.vec3, id: int, @@ -416,6 +416,7 @@ def _ray_hfield( size = hfield_size[hid] adr = hfield_adr[hid] + mat = quat_to_mat(quat) mat_col = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) # compute size and pos of base box @@ -429,16 +430,16 @@ def _ray_hfield( top_pos = pos + mat_col * top_scale # init: intersection with base box - x, _ = _ray_box(base_pos, mat, base_size, pnt, vec) + x, _ = _ray_box(base_pos, quat, base_size, pnt, vec) # check top box: done if no intersection - top_intersect, all = _ray_box(top_pos, mat, top_size, pnt, vec) + top_intersect, all = _ray_box(top_pos, quat, top_size, pnt, vec) if top_intersect < 0.0: return x # map to local frame - lpnt, lvec = _ray_map(pos, mat, pnt, vec) + lpnt, lvec = _ray_map(pos, quat, pnt, vec) # construct basis vectors of normal plane b0 = wp.vec3(1.0, 1.0, 1.0) @@ -543,12 +544,12 @@ def ray_mesh( # In: data_id: int, pos: wp.vec3, - mat: wp.mat33, + quat: wp.quat, pnt: wp.vec3, vec: wp.vec3, ) -> float: """Returns the distance and geomid for ray mesh intersections.""" - pnt, vec = _ray_map(pos, mat, pnt, vec) + pnt, vec = _ray_map(pos, quat, pnt, vec) # compute orthogonal basis vectors if wp.abs(vec[0]) < wp.abs(vec[1]): @@ -601,21 +602,21 @@ def ray_mesh( @wp.func -def ray_geom(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3, geomtype: int) -> float: +def ray_geom(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3, geomtype: int) -> float: """Returns distance along ray to intersection with geom, or infinity if none.""" # TODO(team): static loop unrolling to remove unnecessary branching if geomtype == GeomType.PLANE: - return _ray_plane(pos, mat, size, pnt, vec) + return _ray_plane(pos, quat, size, pnt, vec) elif geomtype == GeomType.SPHERE: return _ray_sphere(pos, size[0] * size[0], pnt, vec) elif geomtype == GeomType.CAPSULE: - return _ray_capsule(pos, mat, size, pnt, vec) + return _ray_capsule(pos, quat, size, pnt, vec) elif geomtype == GeomType.ELLIPSOID: - return _ray_ellipsoid(pos, mat, size, pnt, vec) + return _ray_ellipsoid(pos, quat, size, pnt, vec) elif geomtype == GeomType.CYLINDER: - return _ray_cylinder(pos, mat, size, pnt, vec) + return _ray_cylinder(pos, quat, size, pnt, vec) elif geomtype == GeomType.BOX: - dist, _ = _ray_box(pos, mat, size, pnt, vec) + dist, _ = _ray_box(pos, quat, size, pnt, vec) return dist else: return wp.inf @@ -668,7 +669,7 @@ def _ray_geom_mesh( bodyexclude, ): pos = geom_xpos_in[worldid, geomid] - mat = quat_to_mat(geom_xquat_in[worldid, geomid]) + quat = geom_xquat_in[worldid, geomid] type = geom_type[geomid] if type == GeomType.MESH: @@ -680,7 +681,7 @@ def _ray_geom_mesh( mesh_face, geom_dataid[geomid], pos, - mat, + quat, pnt, vec, ) @@ -694,13 +695,13 @@ def _ray_geom_mesh( hfield_adr, hfield_data, pos, - mat, + quat, pnt, vec, geomid, ) else: - return ray_geom(pos, mat, geom_size[worldid % geom_size.shape[0], geomid], pnt, vec, type) + return ray_geom(pos, quat, geom_size[worldid % geom_size.shape[0], geomid], pnt, vec, type) else: return wp.inf diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index c92b86c49..083294959 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -1442,7 +1442,7 @@ def _accelerometer( site_bodyid: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), cacc_in: wp.array2d(dtype=wp.spatial_vector), @@ -1451,7 +1451,7 @@ def _accelerometer( objid: int, ) -> wp.vec3: bodyid = site_bodyid[objid] - rot = site_xmat_in[worldid, objid] + rot = math.quat_to_mat(site_xquat_in[worldid, objid]) rotT = wp.transpose(rot) cvel = cvel_in[worldid, bodyid] cvel_top = wp.spatial_top(cvel) @@ -1472,7 +1472,7 @@ def _force( # Model: site_bodyid: wp.array(dtype=int), # Data in: - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), # In: worldid: int, @@ -1480,7 +1480,7 @@ def _force( ) -> wp.vec3: bodyid = site_bodyid[objid] cfrc_int = cfrc_int_in[worldid, bodyid] - site_xmat = site_xmat_in[worldid, objid] + site_xmat = math.quat_to_mat(site_xquat_in[worldid, objid]) return wp.transpose(site_xmat) @ wp.spatial_bottom(cfrc_int) @@ -1491,7 +1491,7 @@ def _torque( site_bodyid: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), # In: @@ -1500,7 +1500,7 @@ def _torque( ) -> wp.vec3: bodyid = site_bodyid[objid] cfrc_int = cfrc_int_in[worldid, bodyid] - site_xmat = site_xmat_in[worldid, objid] + site_xmat = math.quat_to_mat(site_xquat_in[worldid, objid]) dif = site_xpos_in[worldid, objid] - subtree_com_in[worldid, body_rootid[bodyid]] return wp.transpose(site_xmat) @ (wp.spatial_top(cfrc_int) - wp.cross(dif, wp.spatial_bottom(cfrc_int))) @@ -1705,7 +1705,7 @@ def _sensor_acc( xipos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), @@ -1948,14 +1948,14 @@ def _sensor_acc( elif sensortype == SensorType.ACCELEROMETER: vec3 = _accelerometer( - body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, cacc_in, worldid, objid + body_rootid, site_bodyid, site_xpos_in, site_xquat_in, subtree_com_in, cvel_in, cacc_in, worldid, objid ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FORCE: - vec3 = _force(site_bodyid, site_xmat_in, cfrc_int_in, worldid, objid) + vec3 = _force(site_bodyid, site_xquat_in, cfrc_int_in, worldid, objid) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.TORQUE: - vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cfrc_int_in, worldid, objid) + vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, site_xquat_in, subtree_com_in, cfrc_int_in, worldid, objid) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.ACTUATORFRC: val = _actuator_force(actuator_force_in, worldid, objid) @@ -2010,7 +2010,7 @@ def _sensor_touch( sensor_touch_adr: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xmat_in: wp.array2d(dtype=wp.mat33), + site_xquat_in: wp.array2d(dtype=wp.quat), contact_pos_in: wp.array(dtype=wp.vec3), contact_frame_in: wp.array(dtype=wp.mat33), contact_dim_in: wp.array(dtype=int), @@ -2066,7 +2066,7 @@ def _sensor_touch( if ( ray.ray_geom( site_xpos_in[worldid, objid], - site_xmat_in[worldid, objid], + site_xquat_in[worldid, objid], site_size[objid], contact_pos_in[conid], conray, @@ -2410,7 +2410,7 @@ def sensor_acc(m: Model, d: Data): m.sensor_adr, m.sensor_touch_adr, d.site_xpos, - d.site_xmat, + d.site_xquat, d.contact.pos, d.contact.frame, d.contact.dim, @@ -2544,7 +2544,7 @@ def sensor_acc(m: Model, d: Data): d.xipos, d.geom_xpos, d.site_xpos, - d.site_xmat, + d.site_xquat, d.cam_xpos, d.subtree_com, d.cvel, From b584a0517e6f268d3371cf64d6bfed01bd8a3609 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:03:36 +0100 Subject: [PATCH 040/107] fix all the tests Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 4 +--- mujoco_warp/_src/smooth_test.py | 8 ++++++-- mujoco_warp/_src/types.py | 2 -- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index b7769cde5..0474e2500 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -200,7 +200,6 @@ def _site_local_to_global( xquat_in: wp.array2d(dtype=wp.quat), # Data out: site_xpos_out: wp.array2d(dtype=wp.vec3), - site_xmat_out: wp.array2d(dtype=wp.mat33), site_xquat_out: wp.array2d(dtype=wp.quat), ): worldid, siteid = wp.tid() @@ -208,7 +207,6 @@ def _site_local_to_global( xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] site_xpos_out[worldid, siteid] = xpos + math.rot_vec_quat(site_pos[worldid % site_pos.shape[0], siteid], xquat) - site_xmat_out[worldid, siteid] = math.quat_to_mat(math.mul_quat(xquat, site_quat[worldid % site_quat.shape[0], siteid])) site_xquat_out[worldid, siteid] = math.mul_quat(xquat, site_quat[worldid % site_quat.shape[0], siteid]) @wp.kernel @@ -321,7 +319,7 @@ def kinematics(m: Model, d: Data): _site_local_to_global, dim=(d.nworld, m.nsite), inputs=[m.site_bodyid, m.site_pos, m.site_quat, d.xpos, d.xquat], - outputs=[d.site_xpos, d.site_xmat, d.site_xquat], + outputs=[d.site_xpos, d.site_xquat], ) diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index 7ee5a4bb8..574e34134 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -125,7 +125,7 @@ def test_kinematics(self): d.geom_xpos, d.geom_xquat, d.site_xpos, - d.site_xmat, + d.site_xquat, ): arr.zero_() @@ -139,6 +139,10 @@ def test_kinematics(self): for i in range(m.ngeom): mujoco.mju_quat2Mat(geom_xmat[i], d.geom_xquat.numpy()[0, i]) + site_xmat = np.zeros((m.nsite, 9), dtype=np.float64) + for i in range(m.nsite): + mujoco.mju_quat2Mat(site_xmat[i], d.site_xquat.numpy()[0, i]) + _assert_eq(d.xanchor.numpy()[0], mjd.xanchor, "xanchor") _assert_eq(d.xaxis.numpy()[0], mjd.xaxis, "xaxis") _assert_eq(d.xpos.numpy()[0], mjd.xpos, "xpos") @@ -148,7 +152,7 @@ def test_kinematics(self): _assert_eq(d.geom_xpos.numpy()[0], mjd.geom_xpos, "geom_xpos") _assert_eq(geom_xmat, mjd.geom_xmat, "geom_xmat") _assert_eq(d.site_xpos.numpy()[0], mjd.site_xpos, "site_xpos") - _assert_eq(d.site_xmat.numpy()[0], mjd.site_xmat.reshape((-1, 3, 3)), "site_xmat") + _assert_eq(site_xmat, mjd.site_xmat, "site_xmat") def test_com_pos(self): """Tests com_pos.""" diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 15c579194..a363e32cb 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1571,7 +1571,6 @@ class Data: xaxis: Cartesian joint axis (nworld, njnt, 3) geom_xpos: Cartesian geom position (nworld, ngeom, 3) site_xpos: Cartesian site position (nworld, nsite, 3) - site_xmat: Cartesian site orientation (nworld, nsite, 3, 3) cam_xpos: Cartesian camera position (nworld, ncam, 3) cam_xmat: Cartesian camera orientation (nworld, ncam, 3, 3) light_xpos: Cartesian light position (nworld, nlight, 3) @@ -1671,7 +1670,6 @@ class Data: xaxis: array("nworld", "njnt", wp.vec3) geom_xpos: array("nworld", "ngeom", wp.vec3) site_xpos: array("nworld", "nsite", wp.vec3) - site_xmat: array("nworld", "nsite", wp.mat33) cam_xpos: array("nworld", "ncam", wp.vec3) cam_xmat: array("nworld", "ncam", wp.mat33) light_xpos: array("nworld", "nlight", wp.vec3) From 1031c180493119ea1016f8c878d317fc6dcd9af9 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:13:45 +0100 Subject: [PATCH 041/107] oobb_filter Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_driver.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index c9a2ee67c..da12cdf8f 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -22,6 +22,7 @@ from .collision_sdf import sdf_narrowphase from .math import upper_tri_index from .math import quat_to_mat +from .math import rot_vec_quat from .types import MJ_MAXVAL from .types import BroadphaseFilter from .types import BroadphaseType @@ -185,8 +186,8 @@ def _obb_filter( margin2: float, xpos1: wp.vec3, xpos2: wp.vec3, - xmat1: wp.mat33, - xmat2: wp.mat33, + xquat1: wp.quat, + xquat2: wp.quat, ) -> bool: """Oriented bounding boxes collision (see Gottschalk et al.), see mj_collideOBB.""" margin = wp.max(margin1, margin2) @@ -197,8 +198,11 @@ def _obb_filter( radius = wp.vec2() # compute centers in local coordinates - xcenter[0] = xmat1 @ center1 + xpos1 - xcenter[1] = xmat2 @ center2 + xpos2 + xcenter[0] = rot_vec_quat(center1, xquat1) + xpos1 + xcenter[1] = rot_vec_quat(center2, xquat2) + xpos2 + + xmat1 = quat_to_mat(xquat1) + xmat2 = quat_to_mat(xquat2) # compute normals in global coordinates normal[0] = wp.vec3(xmat1[0, 0], xmat1[1, 0], xmat1[2, 0]) @@ -260,7 +264,8 @@ def func( margin_id = worldid % ngeom_margin if wp.static(ngeom_margin > 1) else 0 margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2] xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2] - xmat1, xmat2 = quat_to_mat(geom_xquat_in[worldid, geom1]), quat_to_mat(geom_xquat_in[worldid, geom2]) + xquat1, xquat2 = geom_xquat_in[worldid, geom1], geom_xquat_in[worldid, geom2] + xmat1, xmat2 = quat_to_mat(xquat1), quat_to_mat(xquat2) if rbound1 == 0.0 or rbound2 == 0.0: if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE): @@ -273,7 +278,7 @@ def func( if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2): return False if wp.static(opt_broadphase_filter & BroadphaseFilter.OBB): - if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2): + if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xquat1, xquat2): return False return True From e13c2696b10fe33d58e639d459edc9061e62bb3b Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:15:40 +0100 Subject: [PATCH 042/107] aabb_filter Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_driver.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index da12cdf8f..fc0feaede 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -84,16 +84,16 @@ def _aabb_filter( margin2: float, xpos1: wp.vec3, xpos2: wp.vec3, - xmat1: wp.mat33, - xmat2: wp.mat33, + xquat1: wp.quat, + xquat2: wp.quat, ) -> bool: """Axis aligned boxes collision. references: see Ericson, Real-time Collision Detection section 4.2. filterBox: filter contact based on global AABBs. """ - center1 = xmat1 @ center1 + xpos1 - center2 = xmat2 @ center2 + xpos2 + center1 = rot_vec_quat(center1, xquat1) + xpos1 + center2 = rot_vec_quat(center2, xquat2) + xpos2 margin = wp.max(margin1, margin2) @@ -117,10 +117,10 @@ def _aabb_filter( for j in range(2): for k in range(2): corner1 = wp.vec3(sign[i] * size1[0], sign[j] * size1[1], sign[k] * size1[2]) - pos1 = xmat1 @ corner1 + pos1 = rot_vec_quat(corner1, xquat1) corner2 = wp.vec3(sign[i] * size2[0], sign[j] * size2[1], sign[k] * size2[2]) - pos2 = xmat2 @ corner2 + pos2 = rot_vec_quat(corner2, xquat2) if pos1[0] > max_x1: max_x1 = pos1[0] @@ -275,7 +275,7 @@ def func( if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2): return False if wp.static(opt_broadphase_filter & BroadphaseFilter.AABB): - if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2): + if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xquat1, xquat2): return False if wp.static(opt_broadphase_filter & BroadphaseFilter.OBB): if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xquat1, xquat2): From dcebea973119a811f3c5383585b50f31766e2c39 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:17:53 +0100 Subject: [PATCH 043/107] plane_filter Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_driver.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index fc0feaede..e02f62aec 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -50,15 +50,17 @@ def _zero_nacon_ncollision( @wp.func def _plane_filter( - size1: float, size2: float, margin1: float, margin2: float, xpos1: wp.vec3, xpos2: wp.vec3, xmat1: wp.mat33, xmat2: wp.mat33 + size1: float, size2: float, margin1: float, margin2: float, xpos1: wp.vec3, xpos2: wp.vec3, xquat1: wp.quat, xquat2: wp.quat ) -> bool: if size1 == 0.0: # geom1 is a plane - dist = wp.dot(xpos2 - xpos1, wp.vec3(xmat1[0, 2], xmat1[1, 2], xmat1[2, 2])) + mat = quat_to_mat(xquat1) + dist = wp.dot(xpos2 - xpos1, wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2])) return dist <= size2 + wp.max(margin1, margin2) elif size2 == 0.0: # geom2 is a plane - dist = wp.dot(xpos1 - xpos2, wp.vec3(xmat2[0, 2], xmat2[1, 2], xmat2[2, 2])) + mat = quat_to_mat(xquat2) + dist = wp.dot(xpos1 - xpos2, wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2])) return dist <= size1 + wp.max(margin1, margin2) return True @@ -265,11 +267,10 @@ def func( margin1, margin2 = geom_margin[margin_id, geom1], geom_margin[margin_id, geom2] xpos1, xpos2 = geom_xpos_in[worldid, geom1], geom_xpos_in[worldid, geom2] xquat1, xquat2 = geom_xquat_in[worldid, geom1], geom_xquat_in[worldid, geom2] - xmat1, xmat2 = quat_to_mat(xquat1), quat_to_mat(xquat2) if rbound1 == 0.0 or rbound2 == 0.0: if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE): - return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2) + return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xquat1, xquat2) else: if wp.static(opt_broadphase_filter & BroadphaseFilter.SPHERE): if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2): From 4217698aa4a366b658fe645a34dce1c1ff22a6dc Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:23:25 +0100 Subject: [PATCH 044/107] mesh_normals Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index 4387149ff..ea0947fe3 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -22,6 +22,7 @@ from .types import mat43 from .types import mat63 from .math import quat_to_mat +from .math import rot_vec_quat # TODO(team): improve compile time to enable backward pass wp.set_module_options({"enable_backward": False}) @@ -1449,7 +1450,7 @@ def _mesh_normals( # In: feature_dim: int, feature_index: wp.vec3i, - mat: wp.mat33, + rot: wp.quat, vertadr: int, polyadr: int, polynormal: wp.array(dtype=wp.vec3), @@ -1482,7 +1483,7 @@ def _mesh_normals( return 0 # three vertices on mesh define a unique face - normals_out[0] = mat @ polynormal[polyadr + faceset[0]] + normals_out[0] = rot_vec_quat(polynormal[polyadr + faceset[0]], rot) indices_out[0] = faceset[0] return 1 @@ -1498,7 +1499,7 @@ def _mesh_normals( if n == 0: return 0 for i in range(n): - normals_out[i] = mat @ polynormal[polyadr + edgeset[i]] + normals_out[i] = rot_vec_quat(polynormal[polyadr + edgeset[i]], rot) indices_out[i] = edgeset[i] return n @@ -1507,7 +1508,7 @@ def _mesh_normals( v1_num = polymapnum[vertadr + v1] for i in range(v1_num): index = polymap[v1_adr + i] - normals_out[i] = mat @ polynormal[polyadr + index] + normals_out[i] = rot_vec_quat(polynormal[polyadr + index], rot) indices_out[i] = index return v1_num return 0 @@ -1971,7 +1972,7 @@ def multicontact( nnorms1 = _mesh_normals( nface1, feature_index1, - quat_to_mat(geom1.rot), + geom1.rot, geom1.vertadr, geom1.mesh_polyadr, polynormal, @@ -1987,7 +1988,7 @@ def multicontact( nnorms2 = _mesh_normals( nface2, feature_index2, - quat_to_mat(geom2.rot), + geom2.rot, geom2.vertadr, geom2.mesh_polyadr, polynormal, From 80a8c56f3b2837d78acf4051b710d660a904f71e Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:28:24 +0100 Subject: [PATCH 045/107] box_normals Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index ea0947fe3..3dbc1283d 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -1566,7 +1566,7 @@ def _mesh_edge_normals( @wp.func def _box_normals2( # In: - mat: wp.mat33, + quat: wp.quat, n: wp.vec3, # Out: normal_out: wp.array(dtype=wp.vec3), @@ -1576,7 +1576,7 @@ def _box_normals2( face_normals = mat63(1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0) # get local coordinates of the normal - + mat = quat_to_mat(quat) local_n = wp.normalize( wp.vec3( mat[0][0] * n[0] + mat[1][0] * n[1] + mat[2][0] * n[2], @@ -1588,7 +1588,7 @@ def _box_normals2( # determine if there is a side close to the normal for i in range(6): if wp.dot(local_n, face_normals[i]) > FACE_TOL: - normal_out[0] = mat @ face_normals[i] + normal_out[0] = rot_vec_quat(face_normals[i], quat) index_out[0] = i return 1 @@ -1601,7 +1601,7 @@ def _box_normals( # In: feature_dim: int, feature_index: wp.vec3i, - mat: wp.mat33, + quat: wp.quat, dir: wp.vec3, # Out: normal_out: wp.array(dtype=wp.vec3), @@ -1616,7 +1616,7 @@ def _box_normals( x = float((v1 & 1) and (v2 & 1) and (v3 & 1)) - float(not (v1 & 1) and not (v2 & 1) and not (v3 & 1)) y = float((v1 & 2) and (v2 & 2) and (v3 & 2)) - float(not (v1 & 2) and not (v2 & 2) and not (v3 & 2)) z = float((v1 & 4) and (v2 & 4) and (v3 & 4)) - float(not (v1 & 4) and not (v2 & 4) and not (v3 & 4)) - normal_out[0] = mat @ wp.vec3(x, y, z) + normal_out[0] = rot_vec_quat(wp.vec3(x, y, z), quat) sgn = x + y + z if x != 0.0: index_out[c] = 0 @@ -1631,35 +1631,35 @@ def _box_normals( index_out[0] = index_out[0] + 1 if c == 1: return 1 - return _box_normals2(mat, dir, normal_out, index_out) + return _box_normals2(quat, dir, normal_out, index_out) if feature_dim == 2: c = 0 x = float((v1 & 1) and (v2 & 1)) - float(not (v1 & 1) and not (v2 & 1)) y = float((v1 & 2) and (v2 & 2)) - float(not (v1 & 2) and not (v2 & 2)) z = float((v1 & 4) and (v2 & 4)) - float(not (v1 & 4) and not (v2 & 4)) if x != 0.0: - normal_out[c] = mat @ wp.vec3(float(x), 0.0, 0.0) + normal_out[c] = rot_vec_quat(wp.vec3(float(x), 0.0, 0.0), quat) index_out[c] = wp.where(x > 0.0, 0, 1) c += 1 if y != 0.0: - normal_out[c] = mat @ wp.vec3(0.0, y, 0.0) + normal_out[c] = rot_vec_quat(wp.vec3(0.0, y, 0.0), quat) index_out[c] = wp.where(y > 0.0, 2, 3) c += 1 if z != 0.0: - normal_out[c] = mat @ wp.vec3(0.0, 0.0, z) + normal_out[c] = rot_vec_quat(wp.vec3(0.0, 0.0, z), quat) index_out[c] = wp.where(z > 0.0, 4, 5) c += 1 if c == 2: return 2 - return _box_normals2(mat, dir, normal_out, index_out) + return _box_normals2(quat, dir, normal_out, index_out) if feature_dim == 1: x = wp.where(v1 & 1, 1.0, -1.0) y = wp.where(v1 & 2, 1.0, -1.0) z = wp.where(v1 & 4, 1.0, -1.0) - normal_out[0] = mat @ wp.vec3(x, 0.0, 0.0) - normal_out[1] = mat @ wp.vec3(0.0, y, 0.0) - normal_out[2] = mat @ wp.vec3(0.0, 0.0, z) + normal_out[0] = rot_vec_quat(wp.vec3(x, 0.0, 0.0), quat) + normal_out[1] = rot_vec_quat(wp.vec3(0.0, y, 0.0), quat) + normal_out[2] = rot_vec_quat(wp.vec3(0.0, 0.0, z), quat) index_out[0] = wp.where(x > 0.0, 0, 1) index_out[1] = wp.where(y > 0.0, 2, 3) index_out[2] = wp.where(z > 0.0, 4, 5) @@ -1967,7 +1967,7 @@ def multicontact( # get all possible face normals for each geom if geomtype1 == GeomType.BOX: - nnorms1 = _box_normals(nface1, feature_index1, quat_to_mat(geom1.rot), dir_neg, n1, idx1) + nnorms1 = _box_normals(nface1, feature_index1, geom1.rot, dir_neg, n1, idx1) elif geomtype1 == GeomType.MESH: nnorms1 = _mesh_normals( nface1, @@ -1983,7 +1983,7 @@ def multicontact( idx1, ) if geomtype2 == GeomType.BOX: - nnorms2 = _box_normals(nface2, feature_index2, quat_to_mat(geom2.rot), dir, n2, idx2) + nnorms2 = _box_normals(nface2, feature_index2, geom2.rot, dir, n2, idx2) elif geomtype2 == GeomType.MESH: nnorms2 = _mesh_normals( nface2, From 74fb1843d41536a16cb6629d613c7ef8d6537277 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:33:31 +0100 Subject: [PATCH 046/107] box_normals follow up Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index 3dbc1283d..0e709d269 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -23,6 +23,7 @@ from .types import mat63 from .math import quat_to_mat from .math import rot_vec_quat +from .math import quat_inv # TODO(team): improve compile time to enable backward pass wp.set_module_options({"enable_backward": False}) @@ -1576,14 +1577,7 @@ def _box_normals2( face_normals = mat63(1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, -1.0) # get local coordinates of the normal - mat = quat_to_mat(quat) - local_n = wp.normalize( - wp.vec3( - mat[0][0] * n[0] + mat[1][0] * n[1] + mat[2][0] * n[2], - mat[0][1] * n[0] + mat[1][1] * n[1] + mat[2][1] * n[2], - mat[0][2] * n[0] + mat[1][2] * n[1] + mat[2][2] * n[2], - ) - ) + local_n = wp.normalize(rot_vec_quat(n, quat_inv(quat))) # determine if there is a side close to the normal for i in range(6): From baf6e18618e8325b11e4d5c09c18821e58c53555 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:37:24 +0100 Subject: [PATCH 047/107] mesh_edge_normals Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index 0e709d269..43a98e8c6 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -1520,7 +1520,7 @@ def _mesh_normals( def _mesh_edge_normals( # In: dim: int, - mat: wp.mat33, + quat: wp.quat, pos: wp.vec3, vertadr: int, polyadr: int, @@ -1557,7 +1557,7 @@ def _mesh_edge_normals( for j in range(nvert): if polyvert[adr + j] == v1i: k = wp.where(j == 0, nvert - 1, j - 1) - endverts_out[i] = mat @ vert[vertadr + polyvert[adr + k]] + pos + endverts_out[i] = rot_vec_quat(vert[vertadr + polyvert[adr + k]], quat) + pos normals_out[i] = wp.normalize(endverts_out[i] - v1) return v1_num return 0 @@ -2008,7 +2008,7 @@ def multicontact( elif geomtype1 == GeomType.MESH: nnorms1 = _mesh_edge_normals( nface1, - quat_to_mat(geom1.rot), + geom1.rot, geom1.pos, geom1.vertadr, geom1.mesh_polyadr, @@ -2040,7 +2040,7 @@ def multicontact( elif geomtype2 == GeomType.MESH: nnorms2 = _mesh_edge_normals( nface2, - quat_to_mat(geom2.rot), + geom2.rot, geom2.pos, geom2.vertadr, geom2.mesh_polyadr, From 4c8ea653a38e47fd1a484e74f6ad33f646a44767 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:40:52 +0100 Subject: [PATCH 048/107] box_edge_normals Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index 43a98e8c6..246f35faf 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -1666,7 +1666,7 @@ def _box_normals( def _box_edge_normals( # In: dim: int, - mat: wp.mat33, + quat: wp.quat, pos: wp.vec3, size: wp.vec3, v1: wp.vec3, @@ -1687,13 +1687,13 @@ def _box_edge_normals( y = wp.where(v1i & 2, size[1], -size[1]) z = wp.where(v1i & 4, size[2], -size[2]) - endvert_out[0] = mat @ wp.vec3(-x, y, z) + pos + endvert_out[0] = rot_vec_quat(wp.vec3(-x, y, z), quat) + pos normal_out[0] = wp.normalize(endvert_out[0] - v1) - endvert_out[1] = mat @ wp.vec3(x, -y, z) + pos + endvert_out[1] = rot_vec_quat(wp.vec3(x, -y, z), quat) + pos normal_out[1] = wp.normalize(endvert_out[1] - v1) - endvert_out[2] = mat @ wp.vec3(x, y, -z) + pos + endvert_out[2] = rot_vec_quat(wp.vec3(x, y, -z), quat) + pos normal_out[2] = wp.normalize(endvert_out[2] - v1) return 3 return 0 @@ -2003,7 +2003,7 @@ def multicontact( nnorms1 = 0 if geomtype1 == GeomType.BOX: nnorms1 = _box_edge_normals( - nface1, quat_to_mat(geom1.rot), geom1.pos, geom1.size, feature_vertex1[0], feature_vertex1[1], feature_index1[0], n1, endvert + nface1, geom1.rot, geom1.pos, geom1.size, feature_vertex1[0], feature_vertex1[1], feature_index1[0], n1, endvert ) elif geomtype1 == GeomType.MESH: nnorms1 = _mesh_edge_normals( @@ -2035,7 +2035,7 @@ def multicontact( nnorms2 = 0 if geomtype2 == GeomType.BOX: nnorms2 = _box_edge_normals( - nface2, quat_to_mat(geom2.rot), geom2.pos, geom2.size, feature_vertex2[0], feature_vertex2[1], feature_index2[0], n2, endvert + nface2, geom2.rot, geom2.pos, geom2.size, feature_vertex2[0], feature_vertex2[1], feature_index2[0], n2, endvert ) elif geomtype2 == GeomType.MESH: nnorms2 = _mesh_edge_normals( From 91e334a2e12e008d9c1c66cd0846090ba91a4df4 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 14:45:43 +0100 Subject: [PATCH 049/107] box_face Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 54 +++++++++++++++---------------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index 246f35faf..b9a63062d 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -1701,43 +1701,43 @@ def _box_edge_normals( # recover face of a box from its index @wp.func -def _box_face(mat: wp.mat33, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp.array(dtype=wp.vec3)) -> int: +def _box_face(quat: wp.quat, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp.array(dtype=wp.vec3)) -> int: # compute global coordinates of the box face and face normal if idx == 0: # right - face_out[0] = mat @ wp.vec(size[0], size[1], size[2]) + pos - face_out[1] = mat @ wp.vec(size[0], size[1], -size[2]) + pos - face_out[2] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos - face_out[3] = mat @ wp.vec(size[0], -size[1], size[2]) + pos + face_out[0] = rot_vec_quat(wp.vec(size[0], size[1], size[2]), quat) + pos + face_out[1] = rot_vec_quat(wp.vec(size[0], size[1], -size[2]), quat) + pos + face_out[2] = rot_vec_quat(wp.vec(size[0], -size[1], -size[2]), quat) + pos + face_out[3] = rot_vec_quat(wp.vec(size[0], -size[1], size[2]), quat) + pos return 4 if idx == 1: # left - face_out[0] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos - face_out[1] = mat @ wp.vec(-size[0], size[1], size[2]) + pos - face_out[2] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos - face_out[3] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos + face_out[0] = rot_vec_quat(wp.vec(-size[0], size[1], -size[2]), quat) + pos + face_out[1] = rot_vec_quat(wp.vec(-size[0], size[1], size[2]), quat) + pos + face_out[2] = rot_vec_quat(wp.vec(-size[0], -size[1], size[2]), quat) + pos + face_out[3] = rot_vec_quat(wp.vec(-size[0], -size[1], -size[2]), quat) + pos return 4 if idx == 2: # top - face_out[0] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos - face_out[1] = mat @ wp.vec(size[0], size[1], -size[2]) + pos - face_out[2] = mat @ wp.vec(size[0], size[1], size[2]) + pos - face_out[3] = mat @ wp.vec(-size[0], size[1], size[2]) + pos + face_out[0] = rot_vec_quat(wp.vec(-size[0], size[1], -size[2]), quat) + pos + face_out[1] = rot_vec_quat(wp.vec(size[0], size[1], -size[2]), quat) + pos + face_out[2] = rot_vec_quat(wp.vec(size[0], size[1], size[2]), quat) + pos + face_out[3] = rot_vec_quat(wp.vec(-size[0], size[1], size[2]), quat) + pos return 4 if idx == 3: # bottom - face_out[0] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos - face_out[1] = mat @ wp.vec(size[0], -size[1], size[2]) + pos - face_out[2] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos - face_out[3] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos + face_out[0] = rot_vec_quat(wp.vec(-size[0], -size[1], size[2]), quat) + pos + face_out[1] = rot_vec_quat(wp.vec(size[0], -size[1], size[2]), quat) + pos + face_out[2] = rot_vec_quat(wp.vec(size[0], -size[1], -size[2]), quat) + pos + face_out[3] = rot_vec_quat(wp.vec(-size[0], -size[1], -size[2]), quat) + pos return 4 if idx == 4: # front - face_out[0] = mat @ wp.vec(-size[0], size[1], size[2]) + pos - face_out[1] = mat @ wp.vec(size[0], size[1], size[2]) + pos - face_out[2] = mat @ wp.vec(size[0], -size[1], size[2]) + pos - face_out[3] = mat @ wp.vec(-size[0], -size[1], size[2]) + pos + face_out[0] = rot_vec_quat(wp.vec(-size[0], size[1], size[2]), quat) + pos + face_out[1] = rot_vec_quat(wp.vec(size[0], size[1], size[2]), quat) + pos + face_out[2] = rot_vec_quat(wp.vec(size[0], -size[1], size[2]), quat) + pos + face_out[3] = rot_vec_quat(wp.vec(-size[0], -size[1], size[2]), quat) + pos return 4 if idx == 5: # back - face_out[0] = mat @ wp.vec(size[0], size[1], -size[2]) + pos - face_out[1] = mat @ wp.vec(-size[0], size[1], -size[2]) + pos - face_out[2] = mat @ wp.vec(-size[0], -size[1], -size[2]) + pos - face_out[3] = mat @ wp.vec(size[0], -size[1], -size[2]) + pos + face_out[0] = rot_vec_quat(wp.vec(size[0], size[1], -size[2]), quat) + pos + face_out[1] = rot_vec_quat(wp.vec(-size[0], size[1], -size[2]), quat) + pos + face_out[2] = rot_vec_quat(wp.vec(-size[0], -size[1], -size[2]), quat) + pos + face_out[3] = rot_vec_quat(wp.vec(size[0], -size[1], -size[2]), quat) + pos return 4 return 0 @@ -2074,7 +2074,7 @@ def multicontact( else: ind = wp.where(is_edge_contact_geom2, idx1[j], idx1[i]) if geomtype1 == GeomType.BOX: - nface1 = _box_face(quat_to_mat(geom1.rot), geom1.pos, geom1.size, ind, face1) + nface1 = _box_face(geom1.rot, geom1.pos, geom1.size, ind, face1) elif geomtype1 == GeomType.MESH: nface1 = _mesh_face( quat_to_mat(geom1.rot), @@ -2094,7 +2094,7 @@ def multicontact( nface2 = _set_edge(epa_vert2, endvert, face[0], i, face2) else: if geomtype2 == GeomType.BOX: - nface2 = _box_face(quat_to_mat(geom2.rot), geom2.pos, geom2.size, idx2[j], face2) + nface2 = _box_face(geom2.rot, geom2.pos, geom2.size, idx2[j], face2) elif geomtype2 == GeomType.MESH: nface2 = _mesh_face( quat_to_mat(geom2.rot), From e112f6beaca19cea160223c7fff399868929bfb4 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 15:10:00 +0100 Subject: [PATCH 050/107] rotmat Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 41 +++++++++++-------------------- 1 file changed, 15 insertions(+), 26 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index b9a63062d..28e1682f9 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -803,25 +803,14 @@ def _replace_simplex3(pt: Polytope, v1: int, v2: int, v3: int) -> GJKResult: @wp.func -def _rotmat(axis: wp.vec3) -> wp.mat33: - n = wp.norm_l2(axis) - u1 = axis[0] / n - u2 = axis[1] / n - u3 = axis[2] / n - - sin = 0.86602540378 # sin(120 deg) - cos = -0.5 # cos(120 deg) - R = wp.mat33() - R[0, 0] = cos + u1 * u1 * (1.0 - cos) - R[0, 1] = u1 * u2 * (1.0 - cos) - u3 * sin - R[0, 2] = u1 * u3 * (1.0 - cos) + u2 * sin - R[1, 0] = u2 * u1 * (1.0 - cos) + u3 * sin - R[1, 1] = cos + u2 * u2 * (1.0 - cos) - R[1, 2] = u2 * u3 * (1.0 - cos) - u1 * sin - R[2, 0] = u1 * u3 * (1.0 - cos) - u2 * sin - R[2, 1] = u2 * u3 * (1.0 - cos) + u1 * sin - R[2, 2] = cos + u3 * u3 * (1.0 - cos) - return R +def _rotquat(axis: wp.vec3) -> wp.quat: + """Returns quaternion for 120 degree rotation around axis.""" + u = wp.normalize(axis) + # θ = 120°, so θ/2 = 60° + # cos(60°) = 0.5, sin(60°) = √3/2 + half_sin = 0.86602540378 # sin(60°) + half_cos = 0.5 # cos(60°) + return wp.quat(half_cos, half_sin * u[0], half_sin * u[1], half_sin * u[2]) @wp.func @@ -940,9 +929,9 @@ def _polytope2( d1 = wp.cross(e, diff) # rotate around the line segment to get three more points spaced 120 degrees apart - R = _rotmat(diff) - d2 = R @ d1 - d3 = R @ d2 + rot_qaut = _rotquat(diff) + d2 = rot_vec_quat(d1, rot_qaut) + d3 = rot_vec_quat(d2, rot_qaut) # save vertices and get indices for each one pt.vert[0] = simplex[0] @@ -1746,7 +1735,7 @@ def _box_face(quat: wp.quat, pos: wp.vec3, size: wp.vec3, idx: int, face_out: wp @wp.func def _mesh_face( # In: - mat: wp.mat33, + quat: wp.quat, pos: wp.vec3, vertadr: int, polyadr: int, @@ -1763,7 +1752,7 @@ def _mesh_face( nvert = polyvertnum[polyadr + idx] for i in range(nvert - 1, -1, -1): v = vert[vertadr + polyvert[adr + i]] - face_out[j] = mat @ v + pos + face_out[j] = rot_vec_quat(v, quat) + pos j += 1 return nvert @@ -2077,7 +2066,7 @@ def multicontact( nface1 = _box_face(geom1.rot, geom1.pos, geom1.size, ind, face1) elif geomtype1 == GeomType.MESH: nface1 = _mesh_face( - quat_to_mat(geom1.rot), + geom1.rot, geom1.pos, geom1.vertadr, geom1.mesh_polyadr, @@ -2097,7 +2086,7 @@ def multicontact( nface2 = _box_face(geom2.rot, geom2.pos, geom2.size, idx2[j], face2) elif geomtype2 == GeomType.MESH: nface2 = _mesh_face( - quat_to_mat(geom2.rot), + geom2.rot, geom2.pos, geom2.vertadr, geom2.mesh_polyadr, From 2e762d7824658ca9b83e2c022498f43ca0d03bd6 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 15:27:44 +0100 Subject: [PATCH 051/107] wrap Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 4 ++-- mujoco_warp/_src/util_misc.py | 14 +++++------ mujoco_warp/_src/util_misc_test.py | 38 +++++++++++------------------- 3 files changed, 23 insertions(+), 33 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 0474e2500..def0fe18e 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -2719,7 +2719,7 @@ def _spatial_geom_tendon( # get geom information geom_xpos = geom_xpos_in[worldid, wrap_objid_geom] - geom_xmat = math.quat_to_mat(geom_xquat_in[worldid, wrap_objid_geom]) + geom_xquat = geom_xquat_in[worldid, wrap_objid_geom] geomsize = geom_size[worldid, wrap_objid_geom][0] geom_type = wrap_type[wrap_adr] @@ -2736,7 +2736,7 @@ def _spatial_geom_tendon( side = wp.vec3(wp.inf) # compute geom wrap length and connect points (if wrap occurs) - length_geomgeom, geom_pnt0, geom_pnt1 = util_misc.wrap(site_pnt0, site_pnt1, geom_xpos, geom_xmat, geomsize, geom_type, side) + length_geomgeom, geom_pnt0, geom_pnt1 = util_misc.wrap(site_pnt0, site_pnt1, geom_xpos, geom_xquat, geomsize, geom_type, side) # store geom points wrap_geom_xpos_out[worldid, elementid] = wp.spatial_vector(geom_pnt0, geom_pnt1) diff --git a/mujoco_warp/_src/util_misc.py b/mujoco_warp/_src/util_misc.py index 8968795e5..3dda8e620 100644 --- a/mujoco_warp/_src/util_misc.py +++ b/mujoco_warp/_src/util_misc.py @@ -324,7 +324,7 @@ def wrap_inside( @wp.func def wrap( - x0: wp.vec3, x1: wp.vec3, pos: wp.vec3, mat: wp.mat33, radius: float, geomtype: int, side: wp.vec3 + x0: wp.vec3, x1: wp.vec3, pos: wp.vec3, quat: wp.quat, radius: float, geomtype: int, side: wp.vec3 ) -> Tuple[float, wp.vec3, wp.vec3]: """Wrap tendons around spheres and cylinders. @@ -345,9 +345,9 @@ def wrap( return wp.inf, wp.vec3(wp.inf), wp.vec3(wp.inf) # map sites to wrap object's local frame - matT = wp.transpose(mat) - p0 = matT @ (x0 - pos) - p1 = matT @ (x1 - pos) + matT = math.quat_inv(quat) + p0 = math.rot_vec_quat(x0 - pos, matT) + p1 = math.rot_vec_quat(x1 - pos, matT) # too close to origin: return if (wp.norm_l2(p0) < MJ_MINVAL) or (wp.norm_l2(p1) < MJ_MINVAL): @@ -403,7 +403,7 @@ def wrap( if valid_side: # side point: apply same projection as x0, x1 - sidepnt = matT @ (side - pos) + sidepnt = math.rot_vec_quat(side - pos, matT) # side point: project and rescale sidepnt_proj = wp.vec2( @@ -443,8 +443,8 @@ def wrap( wlen = wp.sqrt(wlen * wlen + height * height) # map back to global frame: wpnt - wpnt0 = mat @ res0 + pos - wpnt1 = mat @ res1 + pos + wpnt0 = math.rot_vec_quat(res0, quat) + pos + wpnt1 = math.rot_vec_quat(res1, quat) + pos return wlen, wpnt0, wpnt1 diff --git a/mujoco_warp/_src/util_misc_test.py b/mujoco_warp/_src/util_misc_test.py index 935bef3d0..d3fae1e9c 100644 --- a/mujoco_warp/_src/util_misc_test.py +++ b/mujoco_warp/_src/util_misc_test.py @@ -166,7 +166,7 @@ def _wrap( x0: np.array, x1: np.array, xpos: np.array, - xmat: np.array, + xquat: np.array, radius: float, geomtype: int, side: np.array, @@ -181,7 +181,7 @@ def wrap( x0: wp.vec3, x1: wp.vec3, pos: wp.vec3, - mat: wp.mat33, + quat: wp.quat, radius: float, geomtype: int, side: wp.vec3, @@ -190,7 +190,7 @@ def wrap( wpnt0_out: wp.array(dtype=wp.vec3), wpnt1_out: wp.array(dtype=wp.vec3), ): - length_, wpnt0_, wpnt1_ = util_misc.wrap(x0, x1, pos, mat, radius, geomtype, side) + length_, wpnt0_, wpnt1_ = util_misc.wrap(x0, x1, pos, quat, radius, geomtype, side) length_out[0] = length_ wpnt0_out[0] = wpnt0_ wpnt1_out[0] = wpnt1_ @@ -202,17 +202,7 @@ def wrap( wp.vec3(x0[0], x0[1], x0[2]), wp.vec3(x1[0], x1[1], x1[2]), wp.vec3(xpos[0], xpos[1], xpos[2]), - wp.mat33( - xmat[0, 0], - xmat[0, 1], - xmat[0, 2], - xmat[1, 0], - xmat[1, 1], - xmat[1, 2], - xmat[2, 0], - xmat[2, 1], - xmat[2, 2], - ), + wp.quat(xquat[0], xquat[1], xquat[2], xquat[3]), radius, geomtype, wp.vec3(side[0], side[1], side[2]), @@ -441,11 +431,11 @@ def test_wrap(self, wraptype): x0 = np.array([1, 1, 1]) x1 = np.array([2, 2, 2]) xpos = np.array([0, 0, 0]) - xmat = np.eye(3) + xquat = np.array([1, 0, 0, 0]) radius = 0.1 side = np.array([np.inf, np.inf, np.inf]) - wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side) + wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side) _assert_eq(wlen, -1.0, "wlen") _assert_eq(wpnt0, np.array([np.inf, np.inf, np.inf]), "wpnt0") @@ -455,11 +445,11 @@ def test_wrap(self, wraptype): x0 = np.array([0.1, -1.0, 0.0]) x1 = np.array([0.1, 1.0, 0.0]) xpos = np.array([0, 0, 0]) - xmat = np.eye(3) + xquat = np.array([1, 0, 0, 0]) radius = 0.1 side = np.array([np.inf, np.inf, np.inf]) - wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side) + wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side) _assert_eq(wlen, 0.0, "wlen") _assert_eq(wpnt0, np.array([0.1, 0.0, 0.0]), "wpnt0") @@ -469,11 +459,11 @@ def test_wrap(self, wraptype): x0 = np.array([MJ_MINVAL, -100.0, 0.0]) x1 = np.array([MJ_MINVAL, 100.0, 0.0]) xpos = np.array([0, 0, 0]) - xmat = np.eye(3) + xquat = np.array([1, 0, 0, 0]) radius = 0.1 side = np.array([radius + 10 * MJ_MINVAL, 0, 0]) - wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side) + wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side) # wlen, wpnt0[1], wpnt1[1] are ~0 _assert_eq(wlen, 0, "wlen") @@ -484,12 +474,12 @@ def test_wrap(self, wraptype): x0 = np.array([0.0, -1.0, 0.0]) x1 = np.array([0.0, 1.0, 0.0]) xpos = np.array([0, 0, 0]) - xmat = np.eye(3) + xquat = np.array([1, 0, 0, 0]) radius = 0.1 wraptype = WrapType.CYLINDER side = np.array([0, 0, 0]) - wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side) + wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side) _assert_eq(wlen, -1.0, "wlen") _assert_eq(wpnt0, np.array([np.inf, np.inf, np.inf]), "wpnt0") @@ -499,11 +489,11 @@ def test_wrap(self, wraptype): x0 = np.array([1.0, -1.0, 0.0]) x1 = np.array([1.0, 1.0, 0.0]) xpos = np.array([0, 0, 0]) - xmat = np.eye(3) + xquat = np.array([1, 0, 0, 0]) radius = 0.1 side = np.array([0.0, 0.0, 0.0]) - wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xmat, radius, wraptype, side) + wlen, wpnt0, wpnt1 = _wrap(x0, x1, xpos, xquat, radius, wraptype, side) _assert_eq(wlen, 0.0, "wlen") _assert_eq(wpnt0, np.array([0.1, 0.0, 0.0]), "wpnt0") From 913ec1eac54590e6501531b45752146bce32829c Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 15:37:26 +0100 Subject: [PATCH 052/107] gradient descent Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_sdf.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index 2c65f32fe..73ccb0d8b 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -22,6 +22,9 @@ from .collision_primitive import write_contact from .math import make_frame from .math import quat_to_mat +from .math import rot_vec_quat +from .math import quat_inv +from .math import mul_quat from .ray import ray_mesh from .types import Data from .types import GeomType @@ -582,9 +585,9 @@ def gradient_descent( attr1: wp.vec3, attr2: wp.vec3, pos1: wp.vec3, - rot1: wp.mat33, + rot1: wp.quat, pos2: wp.vec3, - rot2: wp.mat33, + rot2: wp.quat, sdf_type1: int, sdf_type2: int, sdf_iterations: int, @@ -594,8 +597,8 @@ def gradient_descent( mesh_data2: MeshData, ) -> Tuple[float, wp.vec3, wp.vec3]: params = OptimizationParams() - params.rel_mat = wp.transpose(rot1) * rot2 - params.rel_pos = wp.transpose(rot1) * (pos2 - pos1) + params.rel_mat = quat_to_mat(mul_quat(quat_inv(rot1), rot2)) + params.rel_pos = rot_vec_quat(pos2 - pos1, quat_inv(rot1)) params.attr1 = attr1 params.attr2 = attr2 dist, x = gradient_step( @@ -610,8 +613,8 @@ def gradient_descent( grad2 = wp.normalize(grad2) n = grad1 - grad2 n = wp.normalize(n) - pos = rot2 * x + pos2 - n = rot2 * n + pos = rot_vec_quat(x , rot2) + pos2 + n = rot_vec_quat(n, rot2) pos3 = pos - n * dist / 2.0 return dist, pos3, n @@ -768,9 +771,9 @@ def _sdf_narrowphase( aabb_intersection.max = wp.min(aabb1.max, aabb2.max) pos2 = geom2.pos - rot2 = mat_rot_2 + rot2 = geom2.rot pos1 = geom1.pos - rot1 = mat_rot_1 + rot1 = geom1.rot attr1, g1_plugin_id, volume_data1, mesh_data1 = get_sdf_params( oct_child, oct_aabb, oct_coeff, plugin, plugin_attr, type1, geom1.size, g1_plugin, geom_dataid[g1] @@ -810,7 +813,7 @@ def _sdf_narrowphase( aabb_intersection.min[2] + (aabb_intersection.max[2] - aabb_intersection.min[2]) * halton(i, 5), ) x = mat_rot_1 * x_g2 + geom1.pos - x0_initial = wp.transpose(rot2) * (x - pos2) + x0_initial = rot_vec_quat(x - pos2, quat_inv(rot2)) dist, pos, n = gradient_descent( type1, x0_initial, From 3023f7d6c35b9ea5daa3e94daa829201187ff734 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 15:39:54 +0100 Subject: [PATCH 053/107] transform aabb Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_sdf.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index 73ccb0d8b..41e9b6dfb 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -112,7 +112,7 @@ def get_sdf_params( @wp.func -def transform_aabb(aabb_pos: wp.vec3, aabb_size: wp.vec3, pos: wp.vec3, ori: wp.mat33) -> AABB: +def transform_aabb(aabb_pos: wp.vec3, aabb_size: wp.vec3, pos: wp.vec3, ori: wp.quat) -> AABB: aabb = AABB() aabb.max = wp.vec3(-1000000000.0, -1000000000.0, -1000000000.0) aabb.min = wp.vec3(1000000000.0, 1000000000.0, 1000000000.0) @@ -122,7 +122,7 @@ def transform_aabb(aabb_pos: wp.vec3, aabb_size: wp.vec3, pos: wp.vec3, ori: wp. aabb_size.y * (1.0 if (i & 2) else -1.0), aabb_size.z * (1.0 if (i & 4) else -1.0), ) - frame_vec = ori * (vec + aabb_pos) + pos + frame_vec = rot_vec_quat(vec + aabb_pos, ori) + pos aabb.min = wp.min(aabb.min, frame_vec) aabb.max = wp.max(aabb.max, frame_vec) return aabb @@ -756,16 +756,13 @@ def _sdf_narrowphase( g2_plugin = geom_plugin_index[g2] mat_rot_1 = quat_to_mat(geom1.rot) - mat_rot_2 = quat_to_mat(geom2.rot) - g1_to_g2_rot = wp.transpose(mat_rot_1) * mat_rot_2 g1_to_g2_pos = wp.transpose(mat_rot_1) * (geom2.pos - geom1.pos) aabb_pos = geom_aabb[aabb_id, g1, 0] aabb_size = geom_aabb[aabb_id, g1, 1] - identity = wp.identity(3, dtype=float) - aabb1 = transform_aabb(aabb_pos, aabb_size, wp.vec3(0.0), identity) + aabb1 = transform_aabb(aabb_pos, aabb_size, wp.vec3(0.0), wp.quat(1.0, 0.0, 0.0, 0.0)) aabb_pos = geom_aabb[aabb_id, g2, 0] aabb_size = geom_aabb[aabb_id, g2, 1] - aabb2 = transform_aabb(aabb_pos, aabb_size, g1_to_g2_pos, g1_to_g2_rot) + aabb2 = transform_aabb(aabb_pos, aabb_size, g1_to_g2_pos, mul_quat(quat_inv(geom1.rot), geom2.rot)) aabb_intersection = AABB() aabb_intersection.min = wp.max(aabb1.min, aabb2.min) aabb_intersection.max = wp.min(aabb1.max, aabb2.max) From 3aaf2e8efcbc2d04194212391a1f1a4b975ba60e Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 15:42:26 +0100 Subject: [PATCH 054/107] rel_mat Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_sdf.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index 41e9b6dfb..56839a496 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -40,7 +40,7 @@ @wp.struct class OptimizationParams: - rel_mat: wp.mat33 + rel_mat: wp.quat rel_pos: wp.vec3 attr1: wp.vec3 attr2: wp.vec3 @@ -489,7 +489,7 @@ def compute_grad( B = sdf(GeomType.SDF, p2, params.attr2, sdf_type2, volume_data2, mesh_data2) grad1 = sdf_grad(type1, p1, params.attr1, sdf_type1, volume_data1, mesh_data1) grad2 = sdf_grad(GeomType.SDF, p2, params.attr2, sdf_type2, volume_data2, mesh_data2) - grad1_transformed = wp.transpose(params.rel_mat) * grad1 + grad1_transformed = rot_vec_quat(grad1, quat_inv(params.rel_mat)) if sfd_intersection: if A > B: return grad1_transformed @@ -529,7 +529,7 @@ def gradient_step( for i in range(niter): alpha = float(2.0) x2 = wp.vec3(x[0], x[1], x[2]) - x1 = params.rel_mat * x2 + params.rel_pos + x1 = rot_vec_quat(x2, params.rel_mat) + params.rel_pos grad = compute_grad( type1, x1, x2, params, sdf_type1, sdf_type2, sfd_intersection, volume_data1, volume_data2, mesh_data1, mesh_data2 ) @@ -555,7 +555,7 @@ def gradient_step( alpha *= rho wolfe *= rho x = x2 - grad * alpha - x1 = params.rel_mat * x + params.rel_pos + x1 = rot_vec_quat(x, params.rel_mat) + params.rel_pos dist = clearance( type1, x1, @@ -597,7 +597,7 @@ def gradient_descent( mesh_data2: MeshData, ) -> Tuple[float, wp.vec3, wp.vec3]: params = OptimizationParams() - params.rel_mat = quat_to_mat(mul_quat(quat_inv(rot1), rot2)) + params.rel_mat = mul_quat(quat_inv(rot1), rot2) params.rel_pos = rot_vec_quat(pos2 - pos1, quat_inv(rot1)) params.attr1 = attr1 params.attr2 = attr2 @@ -605,9 +605,9 @@ def gradient_descent( type1, x0_initial, params, sdf_type1, sdf_type2, sdf_iterations, False, volume_data1, volume_data2, mesh_data1, mesh_data2 ) dist, x = gradient_step(type1, x, params, sdf_type1, sdf_type2, 1, True, volume_data1, volume_data2, mesh_data1, mesh_data2) - x_1 = params.rel_mat * x + params.rel_pos + x_1 = rot_vec_quat(x, params.rel_mat) + params.rel_pos grad1 = sdf_grad(type1, x_1, params.attr1, sdf_type1, volume_data1, mesh_data1) - grad1 = wp.transpose(params.rel_mat) * grad1 + grad1 = rot_vec_quat(grad1, quat_inv(params.rel_mat)) grad1 = wp.normalize(grad1) grad2 = sdf_grad(GeomType.SDF, x, params.attr2, sdf_type2, volume_data2, mesh_data2) grad2 = wp.normalize(grad2) From f15d020dbeabe8551fb40ba1f1a7fbbdb875f59d Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 15:47:38 +0100 Subject: [PATCH 055/107] hfield_filter Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_convex.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 9c0c2123c..99d723163 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -26,6 +26,9 @@ from .math import make_frame from .math import upper_trid_index from .math import quat_to_mat +from .math import rot_vec_quat +from .math import quat_inv +from .math import mul_quat from .types import MJ_MAX_EPAFACES from .types import MJ_MAX_EPAHORIZON from .types import MJ_MAXCONPAIR @@ -112,10 +115,9 @@ def _hfield_filter( margin_id = worldid % geom_margin.shape[0] pos1 = geom_xpos_in[worldid, g1] - mat1 = quat_to_mat(geom_xquat_in[worldid, g1]) - mat1T = wp.transpose(mat1) pos2 = geom_xpos_in[worldid, g2] - pos = mat1T @ (pos2 - pos1) + rot1_inv = quat_inv(geom_xquat_in[worldid, g1]) + pos = rot_vec_quat(pos2 - pos1, rot1_inv) r2 = geom_rbound[rbound_id, g2] # TODO(team): margin? @@ -133,9 +135,6 @@ def _hfield_filter( if -size1[3] > pos[2] + r2 + margin: # down return True, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf, wp.inf - mat2 = quat_to_mat(geom_xquat_in[worldid, g2]) - mat = mat1T @ mat2 - # aabb for geom in height field frame xmax = -MJ_MAXVAL ymax = -MJ_MAXVAL @@ -148,15 +147,18 @@ def _hfield_filter( center2 = geom_aabb[aabb_id, g2, 0] size2 = geom_aabb[aabb_id, g2, 1] - pos += mat1T @ center2 + pos += rot_vec_quat(center2, rot1_inv) sign = wp.vec2(-1.0, 1.0) + rot2 = geom_xquat_in[worldid, g2] + rot12 = mul_quat(rot1_inv, rot2) + for i in range(2): for j in range(2): for k in range(2): corner_local = wp.vec3(sign[i] * size2[0], sign[j] * size2[1], sign[k] * size2[2]) - corner_hf = mat @ corner_local + corner_hf = rot_vec_quat(corner_local, rot12) if corner_hf[0] > xmax: xmax = corner_hf[0] From 3c3ba067070935063232d4885fc955adb04eafb2 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 16:04:02 +0100 Subject: [PATCH 056/107] convex collision Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_convex.py | 3 +-- mujoco_warp/_src/collision_gjk_test.py | 2 -- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 99d723163..0b4b4f09e 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -25,7 +25,6 @@ from .collision_primitive import write_contact from .math import make_frame from .math import upper_trid_index -from .math import quat_to_mat from .math import rot_vec_quat from .math import quat_inv from .math import mul_quat @@ -644,7 +643,7 @@ def ccd_kernel( x1_ = wp.vec3(0.0, 0.0, 0.0) for i in range(6): x1_ += prism[i] - x1 += quat_to_mat(geom1.rot) @ (x1_ / 6.0) + x1 += rot_vec_quat(x1_ / 6.0, geom1.rot) dist, ncontact, w1, w2, idx = ccd( opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]], diff --git a/mujoco_warp/_src/collision_gjk_test.py b/mujoco_warp/_src/collision_gjk_test.py index 4d6ff9d97..0ebfda547 100644 --- a/mujoco_warp/_src/collision_gjk_test.py +++ b/mujoco_warp/_src/collision_gjk_test.py @@ -30,8 +30,6 @@ from .types import MJ_MAX_EPAHORIZON from .warp_util import nested_kernel -from .math import quat_to_mat - def _geom_dist( m: Model, From 9ebb12973dacaac3a201804a4e5bb59564887828 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 16:36:02 +0100 Subject: [PATCH 057/107] collision_gjk Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index 28e1682f9..41879e31e 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -102,7 +102,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: sp.point = geom.pos + (0.5 * geom.margin) * geom.size[0] * dir return sp - local_dir = wp.transpose(quat_to_mat(geom.rot)) @ dir + local_dir = rot_vec_quat(dir, quat_inv(geom.rot)) if geomtype == GeomType.BOX: tmp = wp.sign(local_dir) res = wp.cw_mul(tmp, geom.size) @@ -114,13 +114,13 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: res = local_dir * geom.size[0] # add cylinder contribution res[2] += wp.sign(local_dir[2]) * geom.size[1] - sp.point = quat_to_mat(geom.rot) @ res + geom.pos + sp.point = rot_vec_quat(res, geom.rot) + geom.pos elif geomtype == GeomType.ELLIPSOID: res = wp.cw_mul(local_dir, geom.size) res = wp.normalize(res) # transform to ellipsoid res = wp.cw_mul(res, geom.size) - sp.point = quat_to_mat(geom.rot) @ res + geom.pos + sp.point = rot_vec_quat(res, geom.rot) + geom.pos elif geomtype == GeomType.CYLINDER: res = wp.vec3(0.0, 0.0, 0.0) # set result in XY plane: support on circle @@ -131,7 +131,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: res[1] = local_dir[1] * scl # set result in Z direction res[2] = wp.sign(local_dir[2]) * geom.size[1] - sp.point = quat_to_mat(geom.rot) @ res + geom.pos + sp.point = rot_vec_quat(res, geom.rot) + geom.pos elif geomtype == GeomType.MESH: max_dist = float(FLOAT_MIN) if geom.graphadr == -1 or geom.vertnum < 10: @@ -173,7 +173,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: sp.vertex_index = geom.graph[vert_globalid + imax] sp.point = geom.vert[geom.vertadr + sp.vertex_index] - sp.point = quat_to_mat(geom.rot) @ sp.point + geom.pos + sp.point = rot_vec_quat(sp.point, geom.rot) + geom.pos elif geomtype == GeomType.HFIELD: max_dist = float(FLOAT_MIN) for i in range(6): @@ -182,7 +182,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: if dist > max_dist: max_dist = dist sp.point = vert - sp.point = quat_to_mat(geom.rot) @ sp.point + geom.pos + sp.point = rot_vec_quat(sp.point, geom.rot) + geom.pos if geom.margin > 0.0: sp.point += dir * (0.5 * geom.margin) From 13c9d3016dc75dd075502b9fa897f4e8e9fbdfae Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 16:46:14 +0100 Subject: [PATCH 058/107] plane normals Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_primitive.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index f0ee8fccf..ba51b7b79 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -33,6 +33,8 @@ from .math import safe_div from .math import upper_trid_index from .math import quat_to_mat +from .math import rot_vec_quat +from .math import quat_inv from .types import MJ_MINMU from .types import MJ_MINVAL from .types import ContactType @@ -112,14 +114,14 @@ def geom_collision_pair( geom1.pos = geom_xpos_in[worldid, g1] geom1.rot = geom_xquat_in[worldid, g1] geom1.size = geom_size[worldid % geom_size.shape[0], g1] - mat1 = quat_to_mat(geom1.rot) - geom1.normal = wp.vec3(mat1[0, 2], mat1[1, 2], mat1[2, 2]) # plane - + if geom_type1 == GeomType.PLANE: + geom1.normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), geom1.rot) + geom2.pos = geom_xpos_in[worldid, g2] geom2.rot = geom_xquat_in[worldid, g2] geom2.size = geom_size[worldid % geom_size.shape[0], g2] - mat2 = quat_to_mat(geom2.rot) - geom2.normal = wp.vec3(mat2[0, 2], mat2[1, 2], mat2[2, 2]) # plane + if geom_type2 == GeomType.PLANE: + geom2.normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), geom2.rot) if geom_type1 == GeomType.MESH: dataid = geom_dataid[g1] From db06ef15a9cdec24d12f4b507979c8836bd1ba22 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 16:54:47 +0100 Subject: [PATCH 059/107] more collision_primitive Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_primitive.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index ba51b7b79..7b4371886 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -189,9 +189,9 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup contact_count = int(0) # get points in the convex frame - mat_transposed = wp.transpose(quat_to_mat(convex.rot)) - plane_pos_local = mat_transposed @ (plane_pos - convex.pos) - n = mat_transposed @ plane_normal + rot_inv = quat_inv(convex.rot) + plane_pos_local = rot_vec_quat(plane_pos - convex.pos, rot_inv) + n = rot_vec_quat(plane_normal, rot_inv) # Store indices in vec4 indices = wp.vec4i(-1, -1, -1, -1) From e059e905b18e28f4790526f8dd1a75acc5dd108c Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 16:56:21 +0100 Subject: [PATCH 060/107] plane_convex Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_primitive.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 7b4371886..eeb9e4fb5 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -386,7 +386,7 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup # Check if the index is unique (appears exactly once) if count == 1: pos = convex.vert[convex.vertadr + idx] - pos = convex.pos + quat_to_mat(convex.rot) @ pos + pos = convex.pos + rot_vec_quat(pos, convex.rot) support = wp.dot(plane_pos_local - convex.vert[convex.vertadr + idx], n) dist = -support pos = pos - 0.5 * dist * plane_normal From 364fc1f37905ce89ab0186d070f062a76f2166b5 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 16:57:31 +0100 Subject: [PATCH 061/107] capsule axis Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_primitive.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index eeb9e4fb5..a7886f8e6 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -736,8 +736,7 @@ def sphere_capsule_wrapper( ): """Calculates one contact between a sphere and a capsule.""" # capsule axis - mat = quat_to_mat(cap.rot) - axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) + axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap.rot) dist, pos, normal = sphere_capsule(sphere.pos, sphere.size[0], cap.pos, axis, cap.size[0], cap.size[1]) From a51dc11fe99577d7b2e16bb9805198ab6a28a244 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 16:58:42 +0100 Subject: [PATCH 062/107] cap-cap Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_primitive.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index a7886f8e6..0813dd58a 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -808,10 +808,8 @@ def capsule_capsule_wrapper( ): """Calculates contacts between two capsules.""" # capsule axes - mat1 = quat_to_mat(cap1.rot) - cap1_axis = wp.vec3(mat1[0, 2], mat1[1, 2], mat1[2, 2]) - mat2 = quat_to_mat(cap2.rot) - cap2_axis = wp.vec3(mat2[0, 2], mat2[1, 2], mat2[2, 2]) + cap1_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap1.rot) + cap2_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap2.rot) dist, pos, normal = capsule_capsule( cap1.pos, From 4581eb05a9f4787838267068ed574bd5cbd32fb9 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:01:21 +0100 Subject: [PATCH 063/107] more axis calcs Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_primitive.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 0813dd58a..803973a24 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -890,8 +890,7 @@ def plane_capsule_wrapper( ): """Calculates contacts between a capsule and a plane.""" # capsule axis - mat = quat_to_mat(cap.rot) - capsule_axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) + capsule_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap.rot) dist, pos, frame = plane_capsule( plane.normal, @@ -1182,8 +1181,7 @@ def sphere_cylinder_wrapper( ): """Calculates contacts between a sphere and a cylinder.""" # cylinder axis - mat = quat_to_mat(cylinder.rot) - cylinder_axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) + cylinder_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cylinder.rot) dist, pos, normal = sphere_cylinder( sphere.pos, From 2de1879576bf2e1606e6a53e1997550b746acbd5 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:02:26 +0100 Subject: [PATCH 064/107] cylinder axis Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_primitive.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 803973a24..1b5799188 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1260,8 +1260,7 @@ def plane_cylinder_wrapper( ): """Calculates contacts between a cylinder and a plane.""" # cylinder axis - mat = quat_to_mat(cylinder.rot) - cylinder_axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) + cylinder_axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cylinder.rot) dist, pos, normal = plane_cylinder( plane.normal, From 2803099c0bd5871c062cabce1698e008326cdb39 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:02:56 +0100 Subject: [PATCH 065/107] one more axis Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_primitive.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 1b5799188..a3bf1e8a5 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1409,8 +1409,7 @@ def capsule_box_wrapper( ): """Calculates contacts between a capsule and a box.""" # Extract capsule axis - mat = quat_to_mat(cap.rot) - axis = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) + axis = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), cap.rot) # Call the core function to get contact geometry dist, pos, normal = capsule_box( From 60bd420798973e1ee2d02707812643bf9996fdbf Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:06:01 +0100 Subject: [PATCH 066/107] more normals Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_driver.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index e02f62aec..b8facf3f2 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -54,13 +54,13 @@ def _plane_filter( ) -> bool: if size1 == 0.0: # geom1 is a plane - mat = quat_to_mat(xquat1) - dist = wp.dot(xpos2 - xpos1, wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2])) + normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), xquat1) + dist = wp.dot(xpos2 - xpos1, normal) return dist <= size2 + wp.max(margin1, margin2) elif size2 == 0.0: # geom2 is a plane - mat = quat_to_mat(xquat2) - dist = wp.dot(xpos1 - xpos2, wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2])) + normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), xquat2) + dist = wp.dot(xpos1 - xpos2, normal) return dist <= size1 + wp.max(margin1, margin2) return True From 36ee7a7136374970ef78375813ae907aec271b29 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:08:49 +0100 Subject: [PATCH 067/107] sdf collisions Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_sdf.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index 56839a496..04cd8fd28 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -21,7 +21,6 @@ from .collision_primitive import geom_collision_pair from .collision_primitive import write_contact from .math import make_frame -from .math import quat_to_mat from .math import rot_vec_quat from .math import quat_inv from .math import mul_quat @@ -755,8 +754,7 @@ def _sdf_narrowphase( g1_plugin = geom_plugin_index[g1] g2_plugin = geom_plugin_index[g2] - mat_rot_1 = quat_to_mat(geom1.rot) - g1_to_g2_pos = wp.transpose(mat_rot_1) * (geom2.pos - geom1.pos) + g1_to_g2_pos = rot_vec_quat(geom2.pos - geom1.pos, quat_inv(geom1.rot)) aabb_pos = geom_aabb[aabb_id, g1, 0] aabb_size = geom_aabb[aabb_id, g1, 1] aabb1 = transform_aabb(aabb_pos, aabb_size, wp.vec3(0.0), wp.quat(1.0, 0.0, 0.0, 0.0)) @@ -809,7 +807,7 @@ def _sdf_narrowphase( aabb_intersection.min[1] + (aabb_intersection.max[1] - aabb_intersection.min[1]) * halton(i, 3), aabb_intersection.min[2] + (aabb_intersection.max[2] - aabb_intersection.min[2]) * halton(i, 5), ) - x = mat_rot_1 * x_g2 + geom1.pos + x = rot_vec_quat(x_g2, geom1.rot) + geom1.pos x0_initial = rot_vec_quat(x - pos2, quat_inv(rot2)) dist, pos, n = gradient_descent( type1, From 83d5a07f47ab2b06f739d8b4e0cd591fae45fa57 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:23:55 +0100 Subject: [PATCH 068/107] fluid Signed-off-by: Alain Denzler --- mujoco_warp/_src/passive.py | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 1896f956f..9b5245a43 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -302,8 +302,6 @@ def _fluid_force( # Body kinematics xipos = xipos_in[worldid, bodyid] - rot = math.quat_to_mat(xiquat_in[worldid, bodyid]) - rotT = wp.transpose(rot) cvel = cvel_in[worldid, bodyid] ang_global = wp.spatial_top(cvel) lin_global = wp.spatial_bottom(cvel) @@ -325,17 +323,17 @@ def _fluid_force( size = geom_size[worldid % geom_size.shape[0], geomid] semiaxes = _geom_semiaxes(size, geom_type[geomid]) - geom_rot = math.quat_to_mat(geom_xquat_in[worldid, geomid]) - geom_rotT = wp.transpose(geom_rot) geom_pos = geom_xpos_in[worldid, geomid] + geom_rot = geom_xquat_in[worldid, geomid] + geom_rot_inv = math.quat_inv(geom_rot) lin_point = lin_com + wp.cross(ang_global, geom_pos - xipos) - l_ang = geom_rotT @ ang_global - l_lin = geom_rotT @ lin_point + l_ang = math.rot_vec_quat(ang_global, geom_rot_inv) + l_lin = math.rot_vec_quat(lin_point, geom_rot_inv) if wind[0] or wind[1] or wind[2]: - l_lin -= geom_rotT @ wind + l_lin -= math.rot_vec_quat(wind, geom_rot_inv) lfrc_torque = wp.vec3(0.0) lfrc_force = wp.vec3(0.0) @@ -431,17 +429,20 @@ def _fluid_force( lfrc_force *= coef # map force/torque from local to world frame: lfrc -> bfrc - torque_global += geom_rot @ lfrc_torque - force_global += geom_rot @ lfrc_force + torque_global += math.rot_vec_quat(lfrc_torque, geom_rot) + force_global += math.rot_vec_quat(lfrc_force, geom_rot) fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force_global, torque_global) return - l_ang = rotT @ ang_global - l_lin = rotT @ lin_com + rot = xiquat_in[worldid, bodyid] + rot_inv = math.quat_inv(rot) + + l_ang = math.rot_vec_quat(ang_global, rot_inv) + l_lin = math.rot_vec_quat(lin_com, rot_inv) if wind[0] or wind[1] or wind[2]: - l_lin -= rotT @ wind + l_lin -= math.rot_vec_quat(wind, rot_inv) lfrc_torque = wp.vec3(0.0) lfrc_force = wp.vec3(0.0) @@ -479,8 +480,8 @@ def _fluid_force( box2 * (box0_pow4 + box1_pow4) * wp.abs(l_ang[2]) * l_ang[2] * scl, ) - torque_global = rot @ lfrc_torque - force_global = rot @ lfrc_force + torque_global = math.rot_vec_quat(lfrc_torque, rot) + force_global = math.rot_vec_quat(lfrc_force, rot) fluid_applied_out[worldid, bodyid] = wp.spatial_vector(force_global, torque_global) From a26d16973bc4cd4670a5d2d0b60cd0c97ebd9ed5 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:33:48 +0100 Subject: [PATCH 069/107] ray Signed-off-by: Alain Denzler --- mujoco_warp/_src/ray.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index c2f30335e..4dd381dcb 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -18,7 +18,8 @@ import warp as wp from .math import safe_div -from .math import quat_to_mat +from .math import quat_inv +from .math import rot_vec_quat from .types import MJ_MINVAL from .types import Data from .types import GeomType @@ -41,13 +42,11 @@ def _ray_map(pos: wp.vec3, quat: wp.quat, pnt: wp.vec3, vec: wp.vec3) -> Tuple[w Returns: 3D point and 3D direction in local geom frame """ - matT = wp.transpose(quat_to_mat(quat)) - lpnt = matT @ (pnt - pos) - lvec = matT @ vec - + quat_inv = quat_inv(quat) + lpnt = rot_vec_quat(pnt - pos, quat_inv) + lvec = rot_vec_quat(vec, quat_inv) return lpnt, lvec - @wp.func def _ray_eliminate( # Model: @@ -416,18 +415,17 @@ def _ray_hfield( size = hfield_size[hid] adr = hfield_adr[hid] - mat = quat_to_mat(quat) - mat_col = wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) + normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), quat) # compute size and pos of base box base_scale = size[3] * 0.5 base_size = wp.vec3(size[0], size[1], base_scale) - base_pos = pos + mat_col * base_scale + base_pos = pos + normal * base_scale # compute size and pos of top box top_scale = size[2] * 0.5 top_size = wp.vec3(size[0], size[1], top_scale) - top_pos = pos + mat_col * top_scale + top_pos = pos + normal * top_scale # init: intersection with base box x, _ = _ray_box(base_pos, quat, base_size, pnt, vec) From ccda19b98399bc77859c08c224196f946d5ca3e3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:38:05 +0100 Subject: [PATCH 070/107] smooth some more Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index def0fe18e..3426e0409 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -1833,8 +1833,7 @@ def _transmission( idslider = trnid[1] gear0 = gear[0] rod = actuator_cranklength[actid] - site_xmat = math.quat_to_mat(site_xquat_in[worldid, idslider]) - axis = wp.vec3(site_xmat[0, 2], site_xmat[1, 2], site_xmat[2, 2]) + axis = math.rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), site_xquat_in[worldid, idslider]) site_xpos_id = site_xpos_in[worldid, id] site_xpos_idslider = site_xpos_in[worldid, idslider] vec = site_xpos_id - site_xpos_idslider @@ -1974,16 +1973,16 @@ def _transmission( site_xpos = site_xpos_in[worldid, siteid] ref_xpos = site_xpos_in[worldid, refid] - ref_xmat = math.quat_to_mat(site_xquat_in[worldid, refid]) + ref_xquat = site_xquat_in[worldid, refid] length = float(0.0) if translational_transmission: # vec: site position in reference site frame - vec = wp.transpose(ref_xmat) @ (site_xpos - ref_xpos) + vec = math.rot_vec_quat(site_xpos - ref_xpos, math.quat_inv(ref_xquat)) length += wp.dot(vec, gear_translation) - wrench_translation = ref_xmat @ gear_translation + wrench_translation = math.rot_vec_quat(gear_translation, ref_xquat) if rotational_transmission: # get site and refsite quats from parent bodies (avoid converting matrix to quat) @@ -1994,7 +1993,7 @@ def _transmission( vec = math.quat_sub(quat, refquat) length += wp.dot(vec, gear_rotational) - wrench_rotation = ref_xmat @ gear_rotational + wrench_rotation = math.rot_vec_quat(gear_rotational, ref_xquat) actuator_length_out[worldid, actid] = length From e20c64d5659bfb500fefca36265a4b3b176323e0 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:47:27 +0100 Subject: [PATCH 071/107] one more conversion Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 3426e0409..0f7d6a008 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -1920,9 +1920,9 @@ def _transmission( # reference site undefined if refid == -1: # wrench: gear expressed in global frame - site_xmat = math.quat_to_mat(site_xquat_in[worldid, siteid]) - wrench_translation = site_xmat @ gear_translation - wrench_rotation = site_xmat @ gear_rotational + site_xquat = site_xquat_in[worldid, siteid] + wrench_translation = math.rot_vec_quat(gear_translation, site_xquat) + wrench_rotation = math.rot_vec_quat(gear_rotational, site_xquat) # moment: global Jacobian projected on wrench # TODO(team): parallelize From 4c65b4394b13fad4ab160b9da67843ebcb012b0a Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Tue, 9 Dec 2025 17:51:24 +0100 Subject: [PATCH 072/107] some sensor stuff Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 083294959..c3436b529 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -118,7 +118,7 @@ def _magnetometer( objid: int, ) -> wp.vec3: magnetic = opt_magnetic[worldid % opt_magnetic.shape[0]] - return wp.transpose(math.quat_to_mat(site_xquat_in[worldid, objid])) @ magnetic + return math.rot_vec_quat(magnetic, math.quat_inv(site_xquat_in[worldid, objid])) @wp.func @@ -204,10 +204,10 @@ def _sensor_rangefinder_init( sensorid = sensor_rangefinder_adr[rfid] objid = sensor_objid[sensorid] site_xpos = site_xpos_in[worldid, objid] - site_xmat = math.quat_to_mat(site_xquat_in[worldid, objid]) + site_normal = math.rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), site_xquat_in[worldid, objid]) pnt_out[worldid, rfid] = site_xpos - vec_out[worldid, rfid] = wp.vec3(site_xmat[0, 2], site_xmat[1, 2], site_xmat[2, 2]) + vec_out[worldid, rfid] = site_normal @wp.func From b81c630d7c7b852fe4b9411bcb29cc17cf690c7e Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 09:42:54 +0100 Subject: [PATCH 073/107] some quaternion math in sensors Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 54 +++++++++++++++++++------------------- 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index c3436b529..927dbd73d 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -313,25 +313,26 @@ def _frame_pos( if reftype == ObjType.BODY: xpos_ref = xipos_in[worldid, refid] - xmat_ref = math.quat_to_mat(xiquat_in[worldid, refid]) + xquat_ref = xiquat_in[worldid, refid] elif objtype == ObjType.XBODY: xpos_ref = xpos_in[worldid, refid] - xmat_ref = math.quat_to_mat(xquat_in[worldid, refid]) + xquat_ref = xquat_in[worldid, refid] elif reftype == ObjType.GEOM: xpos_ref = geom_xpos_in[worldid, refid] - xmat_ref = math.quat_to_mat(geom_xquat_in[worldid, refid]) + xquat_ref = geom_xquat_in[worldid, refid] elif reftype == ObjType.SITE: xpos_ref = site_xpos_in[worldid, refid] - xmat_ref = math.quat_to_mat(site_xquat_in[worldid, refid]) + xquat_ref = site_xquat_in[worldid, refid] elif reftype == ObjType.CAMERA: xpos_ref = cam_xpos_in[worldid, refid] xmat_ref = cam_xmat_in[worldid, refid] + return wp.transpose(xmat_ref) @ (xpos - xpos_ref) else: # UNKNOWN xpos_ref = wp.vec3(0.0) - xmat_ref = wp.identity(3, wp.float32) + xquat_ref = wp.quat(1.0, 0.0, 0.0, 0.0) - return wp.transpose(xmat_ref) @ (xpos - xpos_ref) + return math.rot_vec_quat(xpos - xpos_ref, math.quat_inv(xquat_ref)) @wp.func @@ -348,43 +349,39 @@ def _frame_axis( objtype: int, refid: int, reftype: int, - frame_axis: int, + frame_axis_index: int, + frame_axis_vec: wp.vec3, ) -> wp.vec3: if objtype == ObjType.BODY: - xmat = math.quat_to_mat(xiquat_in[worldid, objid]) - axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) + axis = math.rot_vec_quat(frame_axis_vec, xiquat_in[worldid, objid]) elif objtype == ObjType.XBODY: - xmat = math.quat_to_mat(xquat_in[worldid, objid]) - axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) + axis = math.rot_vec_quat(frame_axis_vec, xquat_in[worldid, objid]) elif objtype == ObjType.GEOM: - xmat = math.quat_to_mat(geom_xquat_in[worldid, objid]) - axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) + axis = math.rot_vec_quat(frame_axis_vec, geom_xquat_in[worldid, objid]) elif objtype == ObjType.SITE: - xmat = math.quat_to_mat(site_xquat_in[worldid, objid]) - axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) + axis = math.rot_vec_quat(frame_axis_vec, site_xquat_in[worldid, objid]) elif objtype == ObjType.CAMERA: xmat = cam_xmat_in[worldid, objid] - axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) - else: # UNKNOWN - axis = wp.vec3(xmat[0, frame_axis], xmat[1, frame_axis], xmat[2, frame_axis]) + axis = wp.vec3(xmat[0, frame_axis_index], xmat[1, frame_axis_index], xmat[2, frame_axis_index]) if refid == -1: return axis if reftype == ObjType.BODY: - xmat_ref = math.quat_to_mat(xiquat_in[worldid, refid]) + xquat_ref = xiquat_in[worldid, refid] elif reftype == ObjType.XBODY: - xmat_ref = math.quat_to_mat(xquat_in[worldid, refid]) + xquat_ref = xquat_in[worldid, refid] elif reftype == ObjType.GEOM: - xmat_ref = math.quat_to_mat(geom_xquat_in[worldid, refid]) + xquat_ref = geom_xquat_in[worldid, refid] elif reftype == ObjType.SITE: - xmat_ref = math.quat_to_mat(site_xquat_in[worldid, refid]) + xquat_ref = site_xquat_in[worldid, refid] elif reftype == ObjType.CAMERA: xmat_ref = cam_xmat_in[worldid, refid] + return wp.transpose(xmat_ref) @ axis else: # UNKNOWN xmat_ref = wp.identity(3, dtype=wp.float32) - return wp.transpose(xmat_ref) @ axis + return math.rot_vec_quat(axis, math.quat_inv(xquat_ref)) @wp.func @@ -565,13 +562,16 @@ def _sensor_pos( refid = sensor_refid[sensorid] reftype = sensor_reftype[sensorid] if sensortype == SensorType.FRAMEXAXIS: - axis = 0 + frame_axis_vec = wp.vec3(1.0, 0.0, 0.0) + frame_axis_index = 0 elif sensortype == SensorType.FRAMEYAXIS: - axis = 1 + frame_axis_vec = wp.vec3(0.0, 1.0, 0.0) + frame_axis_index = 1 elif sensortype == SensorType.FRAMEZAXIS: - axis = 2 + frame_axis_vec = wp.vec3(0.0, 0.0, 1.0) + frame_axis_index = 2 vec3 = _frame_axis( - xquat_in, xiquat_in, geom_xquat_in, site_xquat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis + xquat_in, xiquat_in, geom_xquat_in, site_xquat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, frame_axis_index, frame_axis_vec ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEQUAT: From e8e5d357b9a9e25bb7d7b22eef791a4815a0c569 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 10:04:37 +0100 Subject: [PATCH 074/107] velocimeter Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 927dbd73d..cfaaad842 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -912,13 +912,13 @@ def _velocimeter( ) -> wp.vec3: bodyid = site_bodyid[objid] pos = site_xpos_in[worldid, objid] - rot = math.quat_to_mat(site_xquat_in[worldid, objid]) + rot = site_xquat_in[worldid, objid] cvel = cvel_in[worldid, bodyid] ang = wp.spatial_top(cvel) lin = wp.spatial_bottom(cvel) subtree_com = subtree_com_in[worldid, body_rootid[bodyid]] dif = pos - subtree_com - return wp.transpose(rot) @ (lin - wp.cross(dif, ang)) + return math.rot_vec_quat(lin - wp.cross(dif, ang), math.quat_inv(rot)) @wp.func From caea5b7370842b9a88c05167ddd82889fbfede02 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 10:08:55 +0100 Subject: [PATCH 075/107] frame_linvel Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index cfaaad842..e9dd175aa 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -933,10 +933,10 @@ def _gyro( objid: int, ) -> wp.vec3: bodyid = site_bodyid[objid] - rot = math.quat_to_mat(site_xquat_in[worldid, objid]) + rot = site_xquat_in[worldid, objid] cvel = cvel_in[worldid, bodyid] ang = wp.spatial_top(cvel) - return wp.transpose(rot) @ ang + return math.rot_vec_quat(ang, math.quat_inv(rot)) @wp.func @@ -1082,22 +1082,22 @@ def _frame_linvel( if reftype == ObjType.BODY: xposref = xipos_in[worldid, refid] - xmatref = math.quat_to_mat(xiquat_in[worldid, refid]) + xquatref = xiquat_in[worldid, refid] elif reftype == ObjType.XBODY: xposref = xpos_in[worldid, refid] - xmatref = math.quat_to_mat(xquat_in[worldid, refid]) + xquatref = xquat_in[worldid, refid] elif reftype == ObjType.GEOM: xposref = geom_xpos_in[worldid, refid] - xmatref = math.quat_to_mat(geom_xquat_in[worldid, refid]) + xquatref = geom_xquat_in[worldid, refid] elif reftype == ObjType.SITE: xposref = site_xpos_in[worldid, refid] - xmatref = math.quat_to_mat(site_xquat_in[worldid, refid]) + xquatref = site_xquat_in[worldid, refid] elif reftype == ObjType.CAMERA: xposref = cam_xpos_in[worldid, refid] xmatref = cam_xmat_in[worldid, refid] else: # UNKNOWN xposref = wp.vec3(0.0) - xmatref = wp.identity(3, dtype=float) + xquatref = wp.quat(1.0, 0.0, 0.0, 0.0) cvel, offset = _cvel_offset( body_rootid, @@ -1141,7 +1141,10 @@ def _frame_linvel( xlinvelref = clinvelref - wp.cross(offsetref, cangvelref) rvec = xpos - xposref rel_vel = xlinvel - xlinvelref + wp.cross(rvec, cangvelref) - return wp.transpose(xmatref) @ rel_vel + if reftype == ObjType.CAMERA: + return wp.transpose(xmatref) @ rel_vel + else: + return math.rot_vec_quat(rel_vel, math.quat_inv(xquatref)) else: return xlinvel From b5f2e965be85ff9ab1a3f2d33d8174d86f82ebc6 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 10:10:06 +0100 Subject: [PATCH 076/107] frame_angvel Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index e9dd175aa..515bb899a 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -1196,17 +1196,17 @@ def _frame_angvel( if refid > -1: if reftype == ObjType.BODY: - xmatref = math.quat_to_mat(xiquat_in[worldid, refid]) + xquatref = xiquat_in[worldid, refid] elif reftype == ObjType.XBODY: - xmatref = math.quat_to_mat(xquat_in[worldid, refid]) + xquatref = xquat_in[worldid, refid] elif reftype == ObjType.GEOM: - xmatref = math.quat_to_mat(geom_xquat_in[worldid, refid]) + xquatref = geom_xquat_in[worldid, refid] elif reftype == ObjType.SITE: - xmatref = math.quat_to_mat(site_xquat_in[worldid, refid]) + xquatref = site_xquat_in[worldid, refid] elif reftype == ObjType.CAMERA: xmatref = cam_xmat_in[worldid, refid] else: # UNKNOWN - xmatref = wp.identity(3, dtype=float) + xquatref = wp.quat(1.0, 0.0, 0.0, 0.0) cvelref, _ = _cvel_offset( body_rootid, @@ -1226,7 +1226,10 @@ def _frame_angvel( ) cangvelref = wp.spatial_top(cvelref) - return wp.transpose(xmatref) @ (cangvel - cangvelref) + if reftype == ObjType.CAMERA: + return wp.transpose(xmatref) @ (cangvel - cangvelref) + else: + return math.rot_vec_quat(cangvel - cangvelref, math.quat_inv(xquatref)) else: return cangvel From 8665b7a0c3106cbe7833242f4223fb683a014d29 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 10:13:48 +0100 Subject: [PATCH 077/107] accelerometer Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 515bb899a..0062179e1 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -1457,8 +1457,6 @@ def _accelerometer( objid: int, ) -> wp.vec3: bodyid = site_bodyid[objid] - rot = math.quat_to_mat(site_xquat_in[worldid, objid]) - rotT = wp.transpose(rot) cvel = cvel_in[worldid, bodyid] cvel_top = wp.spatial_top(cvel) cvel_bottom = wp.spatial_bottom(cvel) @@ -1466,9 +1464,10 @@ def _accelerometer( cacc_top = wp.spatial_top(cacc) cacc_bottom = wp.spatial_bottom(cacc) dif = site_xpos_in[worldid, objid] - subtree_com_in[worldid, body_rootid[bodyid]] - ang = rotT @ cvel_top - lin = rotT @ (cvel_bottom - wp.cross(dif, cvel_top)) - acc = rotT @ (cacc_bottom - wp.cross(dif, cacc_top)) + rotT = math.quat_inv(site_xquat_in[worldid, objid]) + ang = math.rot_vec_quat(cvel_top, rotT) + lin = math.rot_vec_quat(cvel_bottom - wp.cross(dif, cvel_top), rotT) + acc = math.rot_vec_quat(cacc_bottom - wp.cross(dif, cacc_top), rotT) correction = wp.cross(ang, lin) return acc + correction From 9577555e73e79945471860f4cbf3545b02af34da Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 10:17:51 +0100 Subject: [PATCH 078/107] force/torque Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 0062179e1..369c32084 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -1485,8 +1485,7 @@ def _force( ) -> wp.vec3: bodyid = site_bodyid[objid] cfrc_int = cfrc_int_in[worldid, bodyid] - site_xmat = math.quat_to_mat(site_xquat_in[worldid, objid]) - return wp.transpose(site_xmat) @ wp.spatial_bottom(cfrc_int) + return math.rot_vec_quat(wp.spatial_bottom(cfrc_int), math.quat_inv(site_xquat_in[worldid, objid])) @wp.func @@ -1505,9 +1504,8 @@ def _torque( ) -> wp.vec3: bodyid = site_bodyid[objid] cfrc_int = cfrc_int_in[worldid, bodyid] - site_xmat = math.quat_to_mat(site_xquat_in[worldid, objid]) dif = site_xpos_in[worldid, objid] - subtree_com_in[worldid, body_rootid[bodyid]] - return wp.transpose(site_xmat) @ (wp.spatial_top(cfrc_int) - wp.cross(dif, wp.spatial_bottom(cfrc_int))) + return math.rot_vec_quat(wp.spatial_top(cfrc_int) - wp.cross(dif, wp.spatial_bottom(cfrc_int)), math.quat_inv(site_xquat_in[worldid, objid])) @wp.func From 4f9d15f18b1b816e019f551536c56be691d7b940 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 10:21:52 +0100 Subject: [PATCH 079/107] final touches Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 0f7d6a008..b8d1daba8 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -2464,11 +2464,11 @@ def _subtree_vel_forward( lin -= wp.cross(xipos - subtree_com_root, ang) subtree_linvel_out[worldid, bodyid] = body_mass[body_mass_id, bodyid] * lin - dv = wp.transpose(ximat) @ ang + dv = math.rot_vec_quat(ang, math.quat_inv(xiquat_in[worldid, bodyid])) dv[0] *= body_inertia[body_inertia_id, bodyid][0] dv[1] *= body_inertia[body_inertia_id, bodyid][1] dv[2] *= body_inertia[body_inertia_id, bodyid][2] - subtree_angmom_out[worldid, bodyid] = ximat @ dv + subtree_angmom_out[worldid, bodyid] = math.rot_vec_quat(dv, xiquat_in[worldid, bodyid]) subtree_bodyvel_out[worldid, bodyid] = wp.spatial_vector(ang, lin) From edf055b0c6505193c9c4831c920f5c1441ad8d3e Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 11:20:47 +0100 Subject: [PATCH 080/107] formatting Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk_test.py | 72 +++++++++++++------------ mujoco_warp/_src/collision_primitive.py | 2 +- mujoco_warp/_src/collision_sdf.py | 2 +- mujoco_warp/_src/io.py | 6 ++- mujoco_warp/_src/io_test.py | 6 +-- mujoco_warp/_src/ray.py | 1 + mujoco_warp/_src/sensor.py | 21 ++++++-- mujoco_warp/_src/smooth.py | 2 + mujoco_warp/_src/types.py | 1 - 9 files changed, 70 insertions(+), 43 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk_test.py b/mujoco_warp/_src/collision_gjk_test.py index 0ebfda547..6d133e787 100644 --- a/mujoco_warp/_src/collision_gjk_test.py +++ b/mujoco_warp/_src/collision_gjk_test.py @@ -558,17 +558,19 @@ def test_cylinder_box(self): ) pos = wp.vec3(0.00015228791744448245, -0.00074981129728257656, 0.29839199781417846680) - rot_mat = np.array([ - 0.99996972084045410156, - 0.00776371126994490623, - -0.00043433305108919740, - -0.00776385562494397163, - 0.99996984004974365234, - -0.00033095158869400620, - 0.00043175052269361913, - 0.00033431366318836808, - 0.99999988079071044922, - ]) + rot_mat = np.array( + [ + 0.99996972084045410156, + 0.00776371126994490623, + -0.00043433305108919740, + -0.00776385562494397163, + 0.99996984004974365234, + -0.00033095158869400620, + 0.00043175052269361913, + 0.00033431366318836808, + 0.99999988079071044922, + ] + ) rot_quat = np.zeros(4) mujoco.mju_mat2Quat(rot_quat, rot_mat) rot = wp.quat(rot_quat[0], rot_quat[1], rot_quat[2], rot_quat[3]) @@ -590,33 +592,37 @@ def test_box_box_float(self): ) pos1 = wp.vec3(-0.17624500393867492676, -0.12375499308109283447, 0.12499777972698211670) - rot1_mat = np.array([ - 1.00000000000000000000, - -0.00000000184385418045, - -0.00000025833372774287, - 0.00000000184391857339, - 1.00000000000000000000, - 0.00000024928382913458, - 0.00000025833372774287, - -0.00000024928382913458, - 1.0000000000000000000, - ]) + rot1_mat = np.array( + [ + 1.00000000000000000000, + -0.00000000184385418045, + -0.00000025833372774287, + 0.00000000184391857339, + 1.00000000000000000000, + 0.00000024928382913458, + 0.00000025833372774287, + -0.00000024928382913458, + 1.0000000000000000000, + ] + ) rot1_quat = np.zeros(4) mujoco.mju_mat2Quat(rot1_quat, rot1_mat) rot1 = wp.quat(rot1_quat[0], rot1_quat[1], rot1_quat[2], rot1_quat[3]) pos2 = wp.vec3(-0.17624500393867492676, -0.12375499308109283447, 0.17499557137489318848) - rot2_mat = np.array([ - 1.00000000000000000000, - -0.00000000184292525685, - 0.00000012980596864054, - 0.00000000184294413064, - 1.00000000000000000000, - -0.00000014602545661546, - -0.00000012980596864054, - 0.00000014602545661546, - 1.00000000000000000000, - ]) + rot2_mat = np.array( + [ + 1.00000000000000000000, + -0.00000000184292525685, + 0.00000012980596864054, + 0.00000000184294413064, + 1.00000000000000000000, + -0.00000014602545661546, + -0.00000012980596864054, + 0.00000014602545661546, + 1.00000000000000000000, + ] + ) rot2_quat = np.zeros(4) mujoco.mju_mat2Quat(rot2_quat, rot2_mat) rot2 = wp.quat(rot2_quat[0], rot2_quat[1], rot2_quat[2], rot2_quat[3]) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index a3bf1e8a5..074d64fab 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -116,7 +116,7 @@ def geom_collision_pair( geom1.size = geom_size[worldid % geom_size.shape[0], g1] if geom_type1 == GeomType.PLANE: geom1.normal = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), geom1.rot) - + geom2.pos = geom_xpos_in[worldid, g2] geom2.rot = geom_xquat_in[worldid, g2] geom2.size = geom_size[worldid % geom_size.shape[0], g2] diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index 04cd8fd28..8f499dbeb 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -612,7 +612,7 @@ def gradient_descent( grad2 = wp.normalize(grad2) n = grad1 - grad2 n = wp.normalize(n) - pos = rot_vec_quat(x , rot2) + pos2 + pos = rot_vec_quat(x, rot2) + pos2 n = rot_vec_quat(n, rot2) pos3 = pos - n * dist / 2.0 return dist, pos3, n diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index f78c80692..b1c890d2b 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -835,7 +835,11 @@ def put_data( d.actuator_moment = wp.array(np.full((nworld, mjm.nu, mjm.nv), actuator_moment), dtype=float) # convert xmat fields to xquat - for xmat, attr, n in [(mjd.ximat, "xiquat", mjm.nbody), (mjd.geom_xmat, "geom_xquat", mjm.ngeom), (mjd.site_xmat, "site_xquat", mjm.nsite)]: + for xmat, attr, n in [ + (mjd.ximat, "xiquat", mjm.nbody), + (mjd.geom_xmat, "geom_xquat", mjm.ngeom), + (mjd.site_xmat, "site_xquat", mjm.nsite), + ]: xquat = np.zeros((n, 4)) for i in range(n): mujoco.mju_mat2Quat(xquat[i], xmat[i]) diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py index d80c66edb..818b329fc 100644 --- a/mujoco_warp/_src/io_test.py +++ b/mujoco_warp/_src/io_test.py @@ -127,13 +127,13 @@ def test_get_data_into(self, nworld, world_id): "xpos", "xquat", "xipos", - #"ximat", + # "ximat", "xanchor", "xaxis", "geom_xpos", - #"geom_xmat", + # "geom_xmat", "site_xpos", - #"site_xmat", + # "site_xmat", "cam_xpos", "cam_xmat", "light_xpos", diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index 4dd381dcb..2f8328cc5 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -47,6 +47,7 @@ def _ray_map(pos: wp.vec3, quat: wp.quat, pnt: wp.vec3, vec: wp.vec3) -> Tuple[w lvec = rot_vec_quat(vec, quat_inv) return lpnt, lvec + @wp.func def _ray_eliminate( # Model: diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 369c32084..de9e07074 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -571,7 +571,18 @@ def _sensor_pos( frame_axis_vec = wp.vec3(0.0, 0.0, 1.0) frame_axis_index = 2 vec3 = _frame_axis( - xquat_in, xiquat_in, geom_xquat_in, site_xquat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, frame_axis_index, frame_axis_vec + xquat_in, + xiquat_in, + geom_xquat_in, + site_xquat_in, + cam_xmat_in, + worldid, + objid, + objtype, + refid, + reftype, + frame_axis_index, + frame_axis_vec, ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEQUAT: @@ -689,7 +700,9 @@ def _sensor_pos( else: return # should not occur refid = sensor_refid[sensorid] - val_bool = inside_geom(site_xpos_in[worldid, refid], site_xquat_in[worldid, refid], site_size[refid], site_type[refid], xpos) + val_bool = inside_geom( + site_xpos_in[worldid, refid], site_xquat_in[worldid, refid], site_size[refid], site_type[refid], xpos + ) _write_scalar(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, float(val_bool), out) elif sensortype == SensorType.E_POTENTIAL: val = energy_in[worldid][0] @@ -1505,7 +1518,9 @@ def _torque( bodyid = site_bodyid[objid] cfrc_int = cfrc_int_in[worldid, bodyid] dif = site_xpos_in[worldid, objid] - subtree_com_in[worldid, body_rootid[bodyid]] - return math.rot_vec_quat(wp.spatial_top(cfrc_int) - wp.cross(dif, wp.spatial_bottom(cfrc_int)), math.quat_inv(site_xquat_in[worldid, objid])) + return math.rot_vec_quat( + wp.spatial_top(cfrc_int) - wp.cross(dif, wp.spatial_bottom(cfrc_int)), math.quat_inv(site_xquat_in[worldid, objid]) + ) @wp.func diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index b8d1daba8..aa8e7298f 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -54,6 +54,7 @@ def _kinematics_root( xipos_out[worldid, 0] = wp.vec3(0.0) xiquat_out[worldid, 0] = wp.quat(1.0, 0.0, 0.0, 0.0) + @wp.kernel def _kinematics_level( # Model: @@ -209,6 +210,7 @@ def _site_local_to_global( site_xpos_out[worldid, siteid] = xpos + math.rot_vec_quat(site_pos[worldid % site_pos.shape[0], siteid], xquat) site_xquat_out[worldid, siteid] = math.mul_quat(xquat, site_quat[worldid % site_quat.shape[0], siteid]) + @wp.kernel def _flex_vertices( # Model: diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index a363e32cb..660b45bb1 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1739,4 +1739,3 @@ class Data: xiquat: array("nworld", "nbody", wp.quat) geom_xquat: array("nworld", "ngeom", wp.quat) site_xquat: array("nworld", "nsite", wp.quat) - From 877ee54cd40824344ef87b92842e4f2f353fd368 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 13:28:48 +0100 Subject: [PATCH 081/107] ruff Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_convex.py | 6 +++--- mujoco_warp/_src/collision_driver.py | 2 +- mujoco_warp/_src/collision_gjk.py | 6 +++--- mujoco_warp/_src/collision_primitive.py | 6 +++--- mujoco_warp/_src/collision_sdf.py | 4 ++-- mujoco_warp/_src/ray.py | 4 ++-- mujoco_warp/_src/util_misc.py | 2 +- 7 files changed, 15 insertions(+), 15 deletions(-) diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 0b4b4f09e..0b730384c 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -24,10 +24,10 @@ from .collision_primitive import geom_collision_pair from .collision_primitive import write_contact from .math import make_frame -from .math import upper_trid_index -from .math import rot_vec_quat -from .math import quat_inv from .math import mul_quat +from .math import quat_inv +from .math import rot_vec_quat +from .math import upper_trid_index from .types import MJ_MAX_EPAFACES from .types import MJ_MAX_EPAHORIZON from .types import MJ_MAXCONPAIR diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index b8facf3f2..4101147fa 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -20,9 +20,9 @@ from .collision_convex import convex_narrowphase from .collision_primitive import primitive_narrowphase from .collision_sdf import sdf_narrowphase -from .math import upper_tri_index from .math import quat_to_mat from .math import rot_vec_quat +from .math import upper_tri_index from .types import MJ_MAXVAL from .types import BroadphaseFilter from .types import BroadphaseType diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index 41879e31e..12792fb77 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -18,12 +18,12 @@ import warp as wp from .collision_primitive import Geom +from .math import quat_inv +from .math import quat_to_mat +from .math import rot_vec_quat from .types import GeomType from .types import mat43 from .types import mat63 -from .math import quat_to_mat -from .math import rot_vec_quat -from .math import quat_inv # TODO(team): improve compile time to enable backward pass wp.set_module_options({"enable_backward": False}) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 074d64fab..f9cc26894 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -30,11 +30,11 @@ from .collision_primitive_core import sphere_cylinder from .collision_primitive_core import sphere_sphere from .math import make_frame -from .math import safe_div -from .math import upper_trid_index +from .math import quat_inv from .math import quat_to_mat from .math import rot_vec_quat -from .math import quat_inv +from .math import safe_div +from .math import upper_trid_index from .types import MJ_MINMU from .types import MJ_MINVAL from .types import ContactType diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index 8f499dbeb..0ff52449a 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -21,9 +21,9 @@ from .collision_primitive import geom_collision_pair from .collision_primitive import write_contact from .math import make_frame -from .math import rot_vec_quat -from .math import quat_inv from .math import mul_quat +from .math import quat_inv +from .math import rot_vec_quat from .ray import ray_mesh from .types import Data from .types import GeomType diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index 2f8328cc5..b1a4d396d 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -17,9 +17,9 @@ import warp as wp -from .math import safe_div from .math import quat_inv from .math import rot_vec_quat +from .math import safe_div from .types import MJ_MINVAL from .types import Data from .types import GeomType @@ -35,7 +35,7 @@ def _ray_map(pos: wp.vec3, quat: wp.quat, pnt: wp.vec3, vec: wp.vec3) -> Tuple[w Args: pos: position of geom frame - mat: orientation of geom frame + quat: orientation of geom frame pnt: starting point of ray in world coordinates vec: direction of ray in world coordinates diff --git a/mujoco_warp/_src/util_misc.py b/mujoco_warp/_src/util_misc.py index 3dda8e620..c5507c66e 100644 --- a/mujoco_warp/_src/util_misc.py +++ b/mujoco_warp/_src/util_misc.py @@ -332,7 +332,7 @@ def wrap( x0: 3D endpoint. x1: 3D endpoint. pos: Position of geom. - mat: Orientation of geom. + quat: Orientation of geom. radius: Geom radius. geomtype: Wrap type (mjtWrap). side: 3D position for sidesite, no side point: wp.vec3(wp.inf). From 84aa371707e92b7036722e70ba3e1f7a3a18dc0a Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 13:34:52 +0100 Subject: [PATCH 082/107] cleanup in io.py Signed-off-by: Alain Denzler --- mujoco_warp/_src/io.py | 41 +++++++++++------------------------------ 1 file changed, 11 insertions(+), 30 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index b1c890d2b..722b3f792 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -914,21 +914,17 @@ def get_data_into( efc_idx = efc_idx[:nefc] # dont emit indices for overflow constraints - # convert xiquat to ximat - ximat = np.zeros((mjm.nbody, 9)) - for i in range(mjm.nbody): - mujoco.mju_quat2Mat(ximat[i], d.xiquat.numpy()[world_id, i]) - result.ximat[:] = ximat - - geom_xmat = np.zeros((mjm.ngeom, 9)) - for i in range(mjm.ngeom): - mujoco.mju_quat2Mat(geom_xmat[i], d.geom_xquat.numpy()[world_id, i]) - result.geom_xmat[:] = geom_xmat - - site_xmat = np.zeros((mjm.nsite, 9)) - for i in range(mjm.nsite): - mujoco.mju_quat2Mat(site_xmat[i], d.site_xquat.numpy()[world_id, i]) - result.site_xmat[:] = site_xmat + # convert quaternions to rotation matrices + for n, quat, xmat in [ + (mjm.nbody, d.xquat, result.xmat), + (mjm.nbody, d.xiquat, result.ximat), + (mjm.ngeom, d.geom_xquat, result.geom_xmat), + (mjm.nsite, d.site_xquat, result.site_xmat), + ]: + mat = np.zeros((n, 9)) + for i in range(n): + mujoco.mju_quat2Mat(mat[i], quat.numpy()[world_id, i]) + xmat[:] = mat result.solver_niter[0] = d.solver_niter.numpy()[world_id] result.ncon = ncon @@ -950,21 +946,6 @@ def get_data_into( result.qacc[:] = d.qacc.numpy()[world_id] result.act_dot[:] = d.act_dot.numpy()[world_id] result.xpos[:] = d.xpos.numpy()[world_id] - xquat = d.xquat.numpy()[world_id] - result.xquat[:] = xquat - # compute xmat from xquat - q = xquat - xmat = np.empty((q.shape[0], 9), dtype=np.float32) - xmat[:, 0] = q[:, 0] ** 2 + q[:, 1] ** 2 - q[:, 2] ** 2 - q[:, 3] ** 2 - xmat[:, 1] = 2.0 * (q[:, 1] * q[:, 2] - q[:, 0] * q[:, 3]) - xmat[:, 2] = 2.0 * (q[:, 1] * q[:, 3] + q[:, 0] * q[:, 2]) - xmat[:, 3] = 2.0 * (q[:, 1] * q[:, 2] + q[:, 0] * q[:, 3]) - xmat[:, 4] = q[:, 0] ** 2 - q[:, 1] ** 2 + q[:, 2] ** 2 - q[:, 3] ** 2 - xmat[:, 5] = 2.0 * (q[:, 2] * q[:, 3] - q[:, 0] * q[:, 1]) - xmat[:, 6] = 2.0 * (q[:, 1] * q[:, 3] - q[:, 0] * q[:, 2]) - xmat[:, 7] = 2.0 * (q[:, 2] * q[:, 3] + q[:, 0] * q[:, 1]) - xmat[:, 8] = q[:, 0] ** 2 - q[:, 1] ** 2 - q[:, 2] ** 2 + q[:, 3] ** 2 - result.xmat[:] = xmat result.xipos[:] = d.xipos.numpy()[world_id] result.xanchor[:] = d.xanchor.numpy()[world_id] result.xaxis[:] = d.xaxis.numpy()[world_id] From 2eb0c0f43c8ee400c1824b6e9be0dc43426512d6 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 13:36:30 +0100 Subject: [PATCH 083/107] fix missing xquat Signed-off-by: Alain Denzler --- mujoco_warp/_src/io.py | 1 + 1 file changed, 1 insertion(+) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 722b3f792..7b364a6d2 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -946,6 +946,7 @@ def get_data_into( result.qacc[:] = d.qacc.numpy()[world_id] result.act_dot[:] = d.act_dot.numpy()[world_id] result.xpos[:] = d.xpos.numpy()[world_id] + result.xquat[:] = d.xquat.numpy()[world_id] result.xipos[:] = d.xipos.numpy()[world_id] result.xanchor[:] = d.xanchor.numpy()[world_id] result.xaxis[:] = d.xaxis.numpy()[world_id] From 89490bd394ef3ea9d4f01b0c5ece7d7ee965b5f4 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 13:40:38 +0100 Subject: [PATCH 084/107] starting to make kernel analyzer happy Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_sdf.py | 2 +- mujoco_warp/_src/passive.py | 4 ++-- mujoco_warp/_src/sensor.py | 14 +++++++------- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index 0ff52449a..3b48c216b 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -665,12 +665,12 @@ def _sdf_narrowphase( geom_plugin_index: wp.array(dtype=int), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xquat_in: wp.array2d(dtype=wp.quat), naconmax_in: int, collision_pair_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array(dtype=wp.vec2i), collision_worldid_in: wp.array(dtype=int), ncollision_in: wp.array(dtype=int), + geom_xquat_in: wp.array2d(dtype=wp.quat), # In: sdf_initpoints: int, sdf_iterations: int, diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 9b5245a43..52b059520 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -280,11 +280,11 @@ def _fluid_force( body_fluid_ellipsoid: wp.array(dtype=bool), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), - xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), + xiquat_in: wp.array2d(dtype=wp.quat), + geom_xquat_in: wp.array2d(dtype=wp.quat), # Out: fluid_applied_out: wp.array2d(dtype=wp.spatial_vector), ): diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index de9e07074..d6ade6387 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -281,13 +281,13 @@ def _frame_pos( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), + xiquat_in: wp.array2d(dtype=wp.quat), + geom_xquat_in: wp.array2d(dtype=wp.quat), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, objid: int, @@ -339,10 +339,10 @@ def _frame_pos( def _frame_axis( # Data in: xquat_in: wp.array2d(dtype=wp.quat), + cam_xmat_in: wp.array2d(dtype=wp.mat33), xiquat_in: wp.array2d(dtype=wp.quat), geom_xquat_in: wp.array2d(dtype=wp.quat), site_xquat_in: wp.array2d(dtype=wp.quat), - cam_xmat_in: wp.array2d(dtype=wp.mat33), # In: worldid: int, objid: int, @@ -489,16 +489,16 @@ def _sensor_pos( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), subtree_com_in: wp.array2d(dtype=wp.vec3), ten_length_in: wp.array2d(dtype=float), actuator_length_in: wp.array2d(dtype=float), + xiquat_in: wp.array2d(dtype=wp.quat), + geom_xquat_in: wp.array2d(dtype=wp.quat), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: rangefinder_dist_in: wp.array2d(dtype=float), sensor_collision_in: wp.array4d(dtype=float), From 2cbd4ab936eb6c18f93ddae79c1053c93f8d4084 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 13:55:20 +0100 Subject: [PATCH 085/107] make analyzer more happy Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_sdf.py | 2 +- mujoco_warp/_src/passive.py | 2 +- mujoco_warp/_src/sensor.py | 85 ++++++++++++++++--------------- mujoco_warp/_src/smooth.py | 20 ++++---- 4 files changed, 55 insertions(+), 54 deletions(-) diff --git a/mujoco_warp/_src/collision_sdf.py b/mujoco_warp/_src/collision_sdf.py index 3b48c216b..17fbf8ade 100644 --- a/mujoco_warp/_src/collision_sdf.py +++ b/mujoco_warp/_src/collision_sdf.py @@ -908,12 +908,12 @@ def sdf_narrowphase(m: Model, d: Data): m.plugin_attr, m.geom_plugin_index, d.geom_xpos, - d.geom_xquat, d.naconmax, d.collision_pair, d.collision_pairid, d.collision_worldid, d.ncollision, + d.geom_xquat, m.opt.sdf_initpoints, m.opt.sdf_iterations, ], diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 52b059520..56fbc64ea 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -508,9 +508,9 @@ def _fluid(m: Model, d: Data): d.xipos, d.xiquat, d.geom_xpos, - d.geom_xquat, d.subtree_com, d.cvel, + d.geom_xquat, ], outputs=[fluid_applied], ) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index d6ade6387..3ae1ef103 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -543,11 +543,8 @@ def _sensor_pos( xpos_in, xquat_in, xipos_in, - xiquat_in, geom_xpos_in, - geom_xquat_in, site_xpos_in, - site_xquat_in, cam_xpos_in, cam_xmat_in, worldid, @@ -555,6 +552,9 @@ def _sensor_pos( objtype, refid, reftype, + xiquat_in, + geom_xquat_in, + site_xquat_in, ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEXAXIS or sensortype == SensorType.FRAMEYAXIS or sensortype == SensorType.FRAMEZAXIS: @@ -572,10 +572,10 @@ def _sensor_pos( frame_axis_index = 2 vec3 = _frame_axis( xquat_in, + cam_xmat_in, xiquat_in, geom_xquat_in, site_xquat_in, - cam_xmat_in, worldid, objid, objtype, @@ -868,11 +868,8 @@ def sensor_pos(m: Model, d: Data): d.xpos, d.xquat, d.xipos, - d.xiquat, d.geom_xpos, - d.geom_xquat, d.site_xpos, - d.site_xquat, d.cam_xpos, d.cam_xmat, d.subtree_com, @@ -880,6 +877,9 @@ def sensor_pos(m: Model, d: Data): d.actuator_length, rangefinder_dist, sensor_collision, + d.xiquat, + d.geom_xquat, + d.site_xquat, ], outputs=[d.sensordata], ) @@ -916,9 +916,9 @@ def _velocimeter( site_bodyid: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, objid: int, @@ -1064,15 +1064,15 @@ def _frame_linvel( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), + xiquat_in: wp.array2d(dtype=wp.quat), + geom_xquat_in: wp.array2d(dtype=wp.quat), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, objid: int, @@ -1173,15 +1173,15 @@ def _frame_angvel( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), + xiquat_in: wp.array2d(dtype=wp.quat), + geom_xquat_in: wp.array2d(dtype=wp.quat), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, objid: int, @@ -1279,11 +1279,8 @@ def _sensor_vel( xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), xipos_in: wp.array2d(dtype=wp.vec3), - xiquat_in: wp.array2d(dtype=wp.quat), geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), cam_xmat_in: wp.array2d(dtype=wp.mat33), subtree_com_in: wp.array2d(dtype=wp.vec3), @@ -1292,6 +1289,9 @@ def _sensor_vel( cvel_in: wp.array2d(dtype=wp.spatial_vector), subtree_linvel_in: wp.array2d(dtype=wp.vec3), subtree_angmom_in: wp.array2d(dtype=wp.vec3), + xiquat_in: wp.array2d(dtype=wp.quat), + geom_xquat_in: wp.array2d(dtype=wp.quat), + site_xquat_in: wp.array2d(dtype=wp.quat), # Data out: sensordata_out: wp.array2d(dtype=float), ): @@ -1302,10 +1302,10 @@ def _sensor_vel( out = sensordata_out[worldid] if sensortype == SensorType.VELOCIMETER: - vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, site_xquat_in, subtree_com_in, cvel_in, worldid, objid) + vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, subtree_com_in, cvel_in, site_xquat_in, worldid, objid) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.GYRO: - vec3 = _gyro(site_bodyid, site_xquat_in, cvel_in, worldid, objid) + vec3 = _gyro(site_bodyid, cvel_in, site_xquat_in, worldid, objid) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.JOINTVEL: val = _joint_vel(jnt_dofadr, qvel_in, worldid, objid) @@ -1331,15 +1331,15 @@ def _sensor_vel( xpos_in, xquat_in, xipos_in, - xiquat_in, geom_xpos_in, - geom_xquat_in, site_xpos_in, - site_xquat_in, cam_xpos_in, cam_xmat_in, subtree_com_in, cvel_in, + xiquat_in, + geom_xquat_in, + site_xquat_in, worldid, objid, objtype, @@ -1359,15 +1359,15 @@ def _sensor_vel( xpos_in, xquat_in, xipos_in, - xiquat_in, geom_xpos_in, - geom_xquat_in, site_xpos_in, - site_xquat_in, cam_xpos_in, cam_xmat_in, subtree_com_in, cvel_in, + xiquat_in, + geom_xquat_in, + site_xquat_in, worldid, objid, objtype, @@ -1414,11 +1414,8 @@ def sensor_vel(m: Model, d: Data): d.xpos, d.xquat, d.xipos, - d.xiquat, d.geom_xpos, - d.geom_xquat, d.site_xpos, - d.site_xquat, d.cam_xpos, d.cam_xmat, d.subtree_com, @@ -1427,6 +1424,9 @@ def sensor_vel(m: Model, d: Data): d.cvel, d.subtree_linvel, d.subtree_angmom, + d.xiquat, + d.geom_xquat, + d.site_xquat, ], outputs=[d.sensordata], ) @@ -1461,10 +1461,10 @@ def _accelerometer( site_bodyid: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), cacc_in: wp.array2d(dtype=wp.spatial_vector), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, objid: int, @@ -1490,8 +1490,8 @@ def _force( # Model: site_bodyid: wp.array(dtype=int), # Data in: - site_xquat_in: wp.array2d(dtype=wp.quat), cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, objid: int, @@ -1508,9 +1508,9 @@ def _torque( site_bodyid: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cfrc_int_in: wp.array2d(dtype=wp.spatial_vector), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, objid: int, @@ -1723,7 +1723,6 @@ def _sensor_acc( xipos_in: wp.array2d(dtype=wp.vec3), geom_xpos_in: wp.array2d(dtype=wp.vec3), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), cam_xpos_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), @@ -1740,6 +1739,7 @@ def _sensor_acc( efc_force_in: wp.array2d(dtype=float), njmax_in: int, nacon_in: wp.array(dtype=int), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: sensor_contact_nmatch_in: wp.array2d(dtype=int), sensor_contact_matchid_in: wp.array3d(dtype=int), @@ -1966,14 +1966,14 @@ def _sensor_acc( elif sensortype == SensorType.ACCELEROMETER: vec3 = _accelerometer( - body_rootid, site_bodyid, site_xpos_in, site_xquat_in, subtree_com_in, cvel_in, cacc_in, worldid, objid + body_rootid, site_bodyid, site_xpos_in, subtree_com_in, cvel_in, cacc_in, site_xquat_in, worldid, objid ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FORCE: - vec3 = _force(site_bodyid, site_xquat_in, cfrc_int_in, worldid, objid) + vec3 = _force(site_bodyid, cfrc_int_in, site_xquat_in, worldid, objid) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.TORQUE: - vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, site_xquat_in, subtree_com_in, cfrc_int_in, worldid, objid) + vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, subtree_com_in, cfrc_int_in, site_xquat_in, worldid, objid) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.ACTUATORFRC: val = _actuator_force(actuator_force_in, worldid, objid) @@ -2028,7 +2028,6 @@ def _sensor_touch( sensor_touch_adr: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), contact_pos_in: wp.array(dtype=wp.vec3), contact_frame_in: wp.array(dtype=wp.mat33), contact_dim_in: wp.array(dtype=int), @@ -2037,6 +2036,7 @@ def _sensor_touch( contact_worldid_in: wp.array(dtype=int), efc_force_in: wp.array2d(dtype=float), nacon_in: wp.array(dtype=int), + site_xquat_in: wp.array2d(dtype=wp.quat), # Data out: sensordata_out: wp.array2d(dtype=float), ): @@ -2128,12 +2128,12 @@ def _sensor_tactile( taxel_sensorid: wp.array(dtype=int), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), contact_geom_in: wp.array(dtype=wp.vec2i), contact_worldid_in: wp.array(dtype=int), nacon_in: wp.array(dtype=int), + geom_xquat_in: wp.array2d(dtype=wp.quat), # Data out: sensordata_out: wp.array2d(dtype=float), ): @@ -2254,7 +2254,6 @@ def _contact_match( sensor_contact_adr: wp.array(dtype=int), # Data in: site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), contact_dist_in: wp.array(dtype=float), contact_pos_in: wp.array(dtype=wp.vec3), contact_frame_in: wp.array(dtype=wp.mat33), @@ -2267,6 +2266,7 @@ def _contact_match( efc_force_in: wp.array2d(dtype=float), njmax_in: int, nacon_in: wp.array(dtype=int), + site_xquat_in: wp.array2d(dtype=wp.quat), # Out: sensor_contact_nmatch_out: wp.array2d(dtype=int), sensor_contact_matchid_out: wp.array3d(dtype=int), @@ -2428,7 +2428,6 @@ def sensor_acc(m: Model, d: Data): m.sensor_adr, m.sensor_touch_adr, d.site_xpos, - d.site_xquat, d.contact.pos, d.contact.frame, d.contact.dim, @@ -2437,6 +2436,7 @@ def sensor_acc(m: Model, d: Data): d.contact.worldid, d.efc.force, d.nacon, + d.site_xquat, ], outputs=[ d.sensordata, @@ -2470,12 +2470,12 @@ def sensor_acc(m: Model, d: Data): m.taxel_vertadr, m.taxel_sensorid, d.geom_xpos, - d.geom_xquat, d.subtree_com, d.cvel, d.contact.geom, d.contact.worldid, d.nacon, + d.geom_xquat, ], outputs=[ d.sensordata, @@ -2509,7 +2509,6 @@ def sensor_acc(m: Model, d: Data): m.sensor_intprm, m.sensor_contact_adr, d.site_xpos, - d.site_xquat, d.contact.dist, d.contact.pos, d.contact.frame, @@ -2522,6 +2521,7 @@ def sensor_acc(m: Model, d: Data): d.efc.force, d.njmax, d.nacon, + d.site_xquat, ], outputs=[sensor_contact_nmatch, sensor_contact_matchid, sensor_contact_criteria, sensor_contact_direction], ) @@ -2562,7 +2562,6 @@ def sensor_acc(m: Model, d: Data): d.xipos, d.geom_xpos, d.site_xpos, - d.site_xquat, d.cam_xpos, d.subtree_com, d.cvel, @@ -2579,6 +2578,8 @@ def sensor_acc(m: Model, d: Data): d.efc.force, d.njmax, d.nacon, + d.site_xquat, + sensor_contact_nmatch, sensor_contact_matchid, sensor_contact_direction, diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index aa8e7298f..0d428b0d2 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -83,9 +83,9 @@ def _kinematics_level( xpos_out: wp.array2d(dtype=wp.vec3), xquat_out: wp.array2d(dtype=wp.quat), xipos_out: wp.array2d(dtype=wp.vec3), - xiquat_out: wp.array2d(dtype=wp.quat), xanchor_out: wp.array2d(dtype=wp.vec3), xaxis_out: wp.array2d(dtype=wp.vec3), + xiquat_out: wp.array2d(dtype=wp.quat), ): worldid, nodeid = wp.tid() bodyid = body_tree_[nodeid] @@ -307,7 +307,7 @@ def kinematics(m: Model, d: Data): d.xquat, body_tree, ], - outputs=[d.xpos, d.xquat, d.xipos, d.xiquat, d.xanchor, d.xaxis], + outputs=[d.xpos, d.xquat, d.xipos, d.xanchor, d.xaxis, d.xiquat], ) wp.launch( @@ -404,8 +404,8 @@ def _cinert( body_inertia: wp.array2d(dtype=wp.vec3), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), - xiquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), + xiquat_in: wp.array2d(dtype=wp.quat), # Data out: cinert_out: wp.array2d(dtype=vec10), ): @@ -510,7 +510,7 @@ def com_pos(m: Model, d: Data): wp.launch( _cinert, dim=(d.nworld, m.nbody), - inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.xiquat, d.subtree_com], + inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.subtree_com, d.xiquat], outputs=[d.cinert], ) wp.launch( @@ -1778,11 +1778,11 @@ def _transmission( qpos_in: wp.array2d(dtype=float), xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), - site_xquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cdof_in: wp.array2d(dtype=wp.spatial_vector), ten_J_in: wp.array3d(dtype=float), ten_length_in: wp.array2d(dtype=float), + site_xquat_in: wp.array2d(dtype=wp.quat), # Data out: actuator_length_out: wp.array2d(dtype=float), actuator_moment_out: wp.array3d(dtype=float), @@ -2180,11 +2180,11 @@ def transmission(m: Model, d: Data): d.qpos, d.xquat, d.site_xpos, - d.site_xquat, d.subtree_com, d.cdof, d.ten_J, d.ten_length, + d.site_xquat, ], outputs=[d.actuator_length, d.actuator_moment], ) @@ -2443,9 +2443,9 @@ def _subtree_vel_forward( body_inertia: wp.array2d(dtype=wp.vec3), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), - xiquat_in: wp.array2d(dtype=wp.quat), subtree_com_in: wp.array2d(dtype=wp.vec3), cvel_in: wp.array2d(dtype=wp.spatial_vector), + xiquat_in: wp.array2d(dtype=wp.quat), # Data out: subtree_linvel_out: wp.array2d(dtype=wp.vec3), subtree_angmom_out: wp.array2d(dtype=wp.vec3), @@ -2557,7 +2557,7 @@ def subtree_vel(m: Model, d: Data): wp.launch( _subtree_vel_forward, dim=(d.nworld, m.nbody), - inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.xiquat, d.subtree_com, d.cvel], + inputs=[m.body_rootid, m.body_mass, m.body_inertia, d.xipos, d.subtree_com, d.cvel, d.xiquat], outputs=[d.subtree_linvel, d.subtree_angmom, d.subtree_bodyvel], ) @@ -2692,10 +2692,10 @@ def _spatial_geom_tendon( wrap_pulley_scale: wp.array(dtype=float), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xquat_in: wp.array2d(dtype=wp.quat), site_xpos_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3), cdof_in: wp.array2d(dtype=wp.spatial_vector), + geom_xquat_in: wp.array2d(dtype=wp.quat), # Data out: ten_J_out: wp.array3d(dtype=float), ten_length_out: wp.array2d(dtype=float), @@ -3050,10 +3050,10 @@ def tendon(m: Model, d: Data): m.wrap_geom_adr, m.wrap_pulley_scale, d.geom_xpos, - d.geom_xquat, d.site_xpos, d.subtree_com, d.cdof, + d.geom_xquat, ], outputs=[d.ten_J, d.ten_length, wrap_geom_xpos], ) From 981d38004b4aea7e353e2da7ac7d4513a71f7ef3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 13:58:00 +0100 Subject: [PATCH 086/107] more fixes Signed-off-by: Alain Denzler --- mujoco_warp/_src/passive.py | 2 +- mujoco_warp/_src/sensor.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 56fbc64ea..1f25e8b23 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -506,10 +506,10 @@ def _fluid(m: Model, d: Data): m.geom_fluid, m.body_fluid_ellipsoid, d.xipos, - d.xiquat, d.geom_xpos, d.subtree_com, d.cvel, + d.xiquat, d.geom_xquat, ], outputs=[fluid_applied], diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 3ae1ef103..c359b97bb 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -875,11 +875,11 @@ def sensor_pos(m: Model, d: Data): d.subtree_com, d.ten_length, d.actuator_length, - rangefinder_dist, - sensor_collision, d.xiquat, d.geom_xquat, d.site_xquat, + rangefinder_dist, + sensor_collision, ], outputs=[d.sensordata], ) From 6e018ab379a30eb6a0d6c9cb9f93bcc7b7f3cb48 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 13:58:50 +0100 Subject: [PATCH 087/107] gyro fix Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index c359b97bb..31872d14a 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -939,8 +939,8 @@ def _gyro( # Model: site_bodyid: wp.array(dtype=int), # Data in: - site_xquat_in: wp.array2d(dtype=wp.quat), cvel_in: wp.array2d(dtype=wp.spatial_vector), + site_xquat_in: wp.array2d(dtype=wp.quat), # In: worldid: int, objid: int, From 6044716ed258cf1f20c49f32059222484edc2671 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 14:00:30 +0100 Subject: [PATCH 088/107] more fixes Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 31872d14a..29ab2b476 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -547,14 +547,14 @@ def _sensor_pos( site_xpos_in, cam_xpos_in, cam_xmat_in, + xiquat_in, + geom_xquat_in, + site_xquat_in, worldid, objid, objtype, refid, - reftype, - xiquat_in, - geom_xquat_in, - site_xquat_in, + reftype ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEXAXIS or sensortype == SensorType.FRAMEYAXIS or sensortype == SensorType.FRAMEZAXIS: From 8337e70e588f0eaef45b3eff83f028befb460754 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 10 Dec 2025 14:00:49 +0100 Subject: [PATCH 089/107] formatting Signed-off-by: Alain Denzler --- mujoco_warp/_src/sensor.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 29ab2b476..9b763e415 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -554,7 +554,7 @@ def _sensor_pos( objid, objtype, refid, - reftype + reftype, ) _write_vector(sensor_type, sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == SensorType.FRAMEXAXIS or sensortype == SensorType.FRAMEYAXIS or sensortype == SensorType.FRAMEZAXIS: @@ -2579,7 +2579,6 @@ def sensor_acc(m: Model, d: Data): d.njmax, d.nacon, d.site_xquat, - sensor_contact_nmatch, sensor_contact_matchid, sensor_contact_direction, From 5d68c79795578529b870af457258833e8e019af6 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 11 Dec 2025 13:34:03 +0100 Subject: [PATCH 090/107] small fix Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index ca25a152c..2ecb568c9 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -897,7 +897,7 @@ def _epa_witness( i3 = pt.vert_index1[face[2]] if geomtype1 == GeomType.HFIELD and (i1 != i2 or i1 != i3): # TODO(kbayes): Fix case where geom2 is near bottom of height field or "extreme" prism heights - n = geom1.rot[:, 2] + n = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), geom1.rot) a = geom1.hfprism[3] b = geom1.hfprism[4] c = geom1.hfprism[5] From aa1300a9cc32a4086396e52e53b3eefbc8f9af30 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 11 Dec 2025 13:46:21 +0100 Subject: [PATCH 091/107] another small fix Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index ad80e5b83..9aa9c956c 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -127,7 +127,7 @@ def test_kinematics(self): ): arr.zero_() - for arr in (d.xquat, d.xmat, d.ximat): + for arr in (d.xquat, d.xiquat): arr_view = arr[:, 1:] arr_view.zero_() From b45b28c8a0d249a1f22a5e0108bd9e0bbb62973f Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 17 Dec 2025 11:15:05 +0100 Subject: [PATCH 092/107] remove init Signed-off-by: Alain Denzler --- mujoco_warp/_src/io.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 577b5bfa3..59920fb5e 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -623,9 +623,6 @@ def make_data( mjd = mujoco.MjData(mjm) mujoco.mj_kinematics(mjm, mjd) - mjd.xquat[0] = (1, 0, 0, 0) - mjd.ximat[0] = np.eye(3) - # mocap mocap_body = np.nonzero(mjm.body_mocapid >= 0)[0] mocap_id = mjm.body_mocapid[mocap_body] From 6b7919166e204b533029651e4c42028aa78621e1 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 17 Dec 2025 11:17:49 +0100 Subject: [PATCH 093/107] remove explicit zeroing, not needed Signed-off-by: Alain Denzler --- mujoco_warp/_src/io.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 59920fb5e..64bd59e81 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -611,12 +611,6 @@ def make_data( contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)}) efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)}) - # world body (body 0): zero position, identity orientation - xquat = np.zeros((nworld, mjm.nbody, 4)) - xquat[:, 0] = (1, 0, 0, 0) - xiquat = np.zeros((nworld, mjm.nbody, 4)) - xiquat[:, 0] = (1, 0, 0, 0) - # world body and static geom (attached to the world) poses are precomputed # this speeds up scenes with many static geoms (e.g. terrains) # TODO(team): remove this when we introduce dof islands + sleeping From e21eaa390559034ccb7f5d6afcad9ac3f492c988 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 17 Dec 2025 11:26:27 +0100 Subject: [PATCH 094/107] fixes Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 2 +- mujoco_warp/_src/io.py | 13 ++++++++++--- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index cadca8024..c868f6ca9 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -2173,7 +2173,7 @@ def _inflate( break if is_side: - n = geom1.rot[:, 2] + n = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), geom1.rot) sp = _support(geom2, geomtype2, x2) x2 = sp.point - margin2 * n diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 64bd59e81..7c6f092ae 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -621,6 +621,14 @@ def make_data( mocap_body = np.nonzero(mjm.body_mocapid >= 0)[0] mocap_id = mjm.body_mocapid[mocap_body] + xiquat = np.zeros((mjm.nbody, 4)) + for i in range(mjm.nbody): + mujoco.mju_mat2Quat(xiquat[i], mjd.ximat[i]) + + geom_xquat = np.zeros((mjm.ngeom, 4)) + for i in range(mjm.ngeom): + mujoco.mju_mat2Quat(geom_xquat[i], mjd.geom_xmat[i]) + d_kwargs = { "qpos": wp.array(np.tile(mjm.qpos0, nworld), shape=(nworld, mjm.nq), dtype=float), "contact": contact, @@ -632,11 +640,10 @@ def make_data( "qLD": None, # world body "xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat), - "xmat": wp.array(np.tile(mjd.xmat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33), - "ximat": wp.array(np.tile(mjd.ximat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.mat33), + "xiquat": wp.array(np.tile(xiquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat), # static geoms "geom_xpos": wp.array(np.tile(mjd.geom_xpos, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.vec3), - "geom_xmat": wp.array(np.tile(mjd.geom_xmat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.mat33), + "geom_xquat": wp.array(np.tile(geom_xquat, (nworld, 1)), shape=(nworld, mjm.ngeom), dtype=wp.quat), # mocap "mocap_pos": wp.array(np.tile(mjm.body_pos[mocap_body[mocap_id]], (nworld, 1)), shape=(nworld, mjm.nmocap), dtype=wp.vec3), "mocap_quat": wp.array( From 8f898e51611d926dff94e2befb48618084ba289c Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 17 Dec 2025 11:30:54 +0100 Subject: [PATCH 095/107] more fixes Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index 79d45a46f..cb98f3b33 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -121,7 +121,7 @@ def test_kinematics(self, make_data): mjd = mujoco.MjData(mjm) d = mjw.make_data(mjm, nworld=nworld) - for arr in (d.xpos, d.xipos, d.xquat, d.xmat, d.ximat, d.xanchor, d.xaxis, d.site_xpos, d.site_xmat): + for arr in (d.xpos, d.xipos, d.xquat, d.xiquat, d.xanchor, d.xaxis, d.site_xpos, d.site_xquat): arr_view = arr[:, 1:] # skip world body arr_view.fill_(wp.inf) From c946f2bfcad89ad88a85264f27ecde249ed4c2ba Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 17 Dec 2025 11:31:57 +0100 Subject: [PATCH 096/107] final fix Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index cb98f3b33..8b0b72319 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -158,11 +158,11 @@ def test_kinematics(self, make_data): _assert_eq(d.xpos.numpy()[i], mjd.xpos, "xpos") _assert_eq(d.xquat.numpy()[i], mjd.xquat, "xquat") _assert_eq(d.xipos.numpy()[i], mjd.xipos, "xipos") - _assert_eq(ximat.numpy()[i], mjd.ximat.reshape((-1, 3, 3)), "ximat") + _assert_eq(ximat[i], mjd.ximat.reshape((-1, 3, 3)), "ximat") _assert_eq(d.geom_xpos.numpy()[i], mjd.geom_xpos, "geom_xpos") - _assert_eq(geom_xmat.numpy()[i], mjd.geom_xmat.reshape((-1, 3, 3)), "geom_xmat") + _assert_eq(geom_xmat[i], mjd.geom_xmat.reshape((-1, 3, 3)), "geom_xmat") _assert_eq(d.site_xpos.numpy()[i], mjd.site_xpos, "site_xpos") - _assert_eq(site_xmat.numpy()[i], mjd.site_xmat.reshape((-1, 3, 3)), "site_xmat") + _assert_eq(site_xmat[i], mjd.site_xmat.reshape((-1, 3, 3)), "site_xmat") _assert_eq(d.mocap_pos.numpy()[i], mjd.mocap_pos, "mocap_pos") _assert_eq(d.mocap_quat.numpy()[i], mjd.mocap_quat, "mocap_quat") From ee488f83ef5b350dfdbc7238d71d12c0a0508b09 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 17 Dec 2025 11:32:54 +0100 Subject: [PATCH 097/107] more test fix Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index 8b0b72319..6f0be28c1 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -158,11 +158,11 @@ def test_kinematics(self, make_data): _assert_eq(d.xpos.numpy()[i], mjd.xpos, "xpos") _assert_eq(d.xquat.numpy()[i], mjd.xquat, "xquat") _assert_eq(d.xipos.numpy()[i], mjd.xipos, "xipos") - _assert_eq(ximat[i], mjd.ximat.reshape((-1, 3, 3)), "ximat") + _assert_eq(ximat[i], mjd.ximat, "ximat") _assert_eq(d.geom_xpos.numpy()[i], mjd.geom_xpos, "geom_xpos") - _assert_eq(geom_xmat[i], mjd.geom_xmat.reshape((-1, 3, 3)), "geom_xmat") + _assert_eq(geom_xmat[i], mjd.geom_xmat, "geom_xmat") _assert_eq(d.site_xpos.numpy()[i], mjd.site_xpos, "site_xpos") - _assert_eq(site_xmat[i], mjd.site_xmat.reshape((-1, 3, 3)), "site_xmat") + _assert_eq(site_xmat[i], mjd.site_xmat, "site_xmat") _assert_eq(d.mocap_pos.numpy()[i], mjd.mocap_pos, "mocap_pos") _assert_eq(d.mocap_quat.numpy()[i], mjd.mocap_quat, "mocap_quat") From 57ebf9f5b86772507c681eaf50a1aa7dc4d6629a Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 17 Dec 2025 11:38:28 +0100 Subject: [PATCH 098/107] nworld Signed-off-by: Alain Denzler --- mujoco_warp/_src/smooth_test.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/mujoco_warp/_src/smooth_test.py b/mujoco_warp/_src/smooth_test.py index 6f0be28c1..a9b1a0aeb 100644 --- a/mujoco_warp/_src/smooth_test.py +++ b/mujoco_warp/_src/smooth_test.py @@ -140,17 +140,20 @@ def test_kinematics(self, make_data): mujoco.mj_kinematics(mjm, mjd) mjw.kinematics(m, d) - ximat = np.zeros((m.nbody, 9), dtype=np.float64) - for i in range(m.nbody): - mujoco.mju_quat2Mat(ximat[i], d.xiquat.numpy()[0, i]) + ximat = np.zeros((nworld, m.nbody, 9), dtype=np.float64) + for i in range(nworld): + for j in range(m.nbody): + mujoco.mju_quat2Mat(ximat[i, j], d.xiquat.numpy()[i, j]) - geom_xmat = np.zeros((m.ngeom, 9), dtype=np.float64) - for i in range(m.ngeom): - mujoco.mju_quat2Mat(geom_xmat[i], d.geom_xquat.numpy()[0, i]) + geom_xmat = np.zeros((nworld, m.ngeom, 9), dtype=np.float64) + for i in range(nworld): + for j in range(m.ngeom): + mujoco.mju_quat2Mat(geom_xmat[i, j], d.geom_xquat.numpy()[i, j]) - site_xmat = np.zeros((m.nsite, 9), dtype=np.float64) - for i in range(m.nsite): - mujoco.mju_quat2Mat(site_xmat[i], d.site_xquat.numpy()[0, i]) + site_xmat = np.zeros((nworld, m.nsite, 9), dtype=np.float64) + for i in range(nworld): + for j in range(m.nsite): + mujoco.mju_quat2Mat(site_xmat[i, j], d.site_xquat.numpy()[i, j]) for i in range(nworld): _assert_eq(d.xanchor.numpy()[i], mjd.xanchor, "xanchor") From 5fef4e983ee602fa0ade1dcdf349cce77bc17608 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 5 Jan 2026 10:24:36 +0100 Subject: [PATCH 099/107] updates Signed-off-by: Alain Denzler --- mujoco_warp/_src/ray.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index affc9d1ea..6fc21bbf2 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -185,7 +185,7 @@ def _ray_triangle( @wp.func def _ray_plane(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]: - """Returns the distance at which a ray intersects with a plane.""" + """Returns the distance and normal at which a ray intersects with a plane.""" # map to local frame lpnt, lvec = _ray_map(pos, quat, pnt, vec) @@ -202,7 +202,7 @@ def _ray_plane(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp # accept only within rendered rectangle if (size[0] <= 0.0 or wp.abs(p[0]) <= size[0]) and (size[1] <= 0.0 or wp.abs(p[1]) <= size[1]): - return x, wp.vec3(mat[0, 2], mat[1, 2], mat[2, 2]) + return x, rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), quat) else: return -1.0, wp.vec3() @@ -226,7 +226,7 @@ def _ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> Tu @wp.func def _ray_capsule(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]: - """Returns the distance at which a ray intersects with a capsule.""" + """Returns the distance and normal at which a ray intersects with a capsule.""" # bounding sphere test ssz = size[0] + size[1] dist_sphere, normal_sphere = _ray_sphere(pos, ssz * ssz, pnt, vec) @@ -292,14 +292,14 @@ def _ray_capsule(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: # normalize, rotate into global frame normal = wp.normalize(normal) - normal = mat @ normal + normal = rot_vec_quat(normal, quat) return x, normal @wp.func def _ray_ellipsoid(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]: - """Returns the distance at which a ray intersects with an ellipsoid.""" + """Returns the distance and normal at which a ray intersects with an ellipsoid.""" # map to local frame lpnt, lvec = _ray_map(pos, quat, pnt, vec) @@ -323,14 +323,14 @@ def _ray_ellipsoid(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec # gradient of ellipsoid function normal = wp.cw_mul(s, l) normal = wp.normalize(normal) - normal = mat @ normal + normal = rot_vec_quat(normal, quat) return sol, normal @wp.func def _ray_cylinder(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]: - """Returns the distance at which a ray intersects with a cylinder.""" + """Returns the distance and normal at which a ray intersects with a cylinder.""" # bounding sphere test ssz = size[0] * size[0] + size[1] * size[1] dist_sphere, normal_sphere = _ray_sphere(pos, ssz, pnt, vec) @@ -384,7 +384,7 @@ def _ray_cylinder(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: else: normal = wp.vec3(0.0, 0.0, float(part)) - normal = mat @ normal + normal = rot_vec_quat(normal, quat) return x, normal @@ -394,7 +394,7 @@ def _ray_cylinder(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: @wp.func def _ray_box(pos: wp.vec3, quat: wp.quat, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6, wp.vec3]: - """Returns the distance at which a ray intersects with a box.""" + """Returns the distance, per side information, and normal at which a ray intersects with a box.""" all = vec6(-1.0, -1.0, -1.0, -1.0, -1.0, -1.0) # bounding sphere test @@ -631,7 +631,7 @@ def ray_mesh( pnt: wp.vec3, vec: wp.vec3, ) -> Tuple[float, wp.vec3]: - """Returns the distance and geomid for ray mesh intersections.""" + """Returns the distance and normal for ray mesh intersections.""" pnt, vec = _ray_map(pos, quat, pnt, vec) # compute orthogonal basis vectors @@ -683,7 +683,7 @@ def ray_mesh( x = dist normal = normal_tri - normal = mat @ normal + normal = rot_vec_quat(normal, quat) return x, normal From 4f7fb9d620722e99c7f1c567591e21e6d6a06d96 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 5 Jan 2026 10:27:15 +0100 Subject: [PATCH 100/107] missing upadte Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_convex.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 1a49c7028..a3df4b021 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -249,7 +249,7 @@ def ccd_hfield_kernel( # Data in: naconmax_in: int, geom_xpos_in: wp.array2d(dtype=wp.vec3), - geom_xmat_in: wp.array2d(dtype=wp.mat33), + geom_xquat_in: wp.array2d(dtype=wp.quat), collision_pair_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array(dtype=wp.vec2i), collision_worldid_in: wp.array(dtype=int), From b41b334dd05f74ea4cc2fedadcfaf1f3b7e8c8ba Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 5 Jan 2026 10:28:36 +0100 Subject: [PATCH 101/107] another missing change. Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_convex.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index a3df4b021..5619989b4 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -297,7 +297,7 @@ def ccd_hfield_kernel( # height field filter no_hf_collision, xmin, xmax, ymin, ymax, zmin, zmax = _hfield_filter( - geom_dataid, geom_aabb, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xmat_in, worldid, g1, g2 + geom_dataid, geom_aabb, geom_rbound, geom_margin, hfield_size, geom_xpos_in, geom_xquat_in, worldid, g1, g2 ) if no_hf_collision: return @@ -343,7 +343,7 @@ def ccd_hfield_kernel( mesh_polymapnum, mesh_polymap, geom_xpos_in, - geom_xmat_in, + geom_xquat_in, geoms, worldid, ) @@ -1339,7 +1339,7 @@ def _pair_count(p1: int, p2: int) -> int: m.pair_friction, d.naconmax, d.geom_xpos, - d.geom_xmat, + d.geom_xquat, d.collision_pair, d.collision_pairid, d.collision_worldid, From c35e2d428eed1f0419963252867f548e08f33fa6 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 5 Jan 2026 10:51:16 +0100 Subject: [PATCH 102/107] missing updates Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_convex.py | 2 +- mujoco_warp/_src/ray.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 5619989b4..7a781850d 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -450,7 +450,7 @@ def ccd_hfield_kernel( x1_ = wp.vec3(0.0, 0.0, 0.0) for i in range(6): x1_ += prism[i] - x1 += geom1.rot @ (x1_ / 6.0) + x1 += rot_vec_quat(x1_ / 6.0, geom1.rot) dist, ncontact, w1, w2, idx = ccd( opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]], diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index 6fc21bbf2..860e6c75a 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -475,20 +475,20 @@ def _ray_hfield( size = hfield_size[hid] adr = hfield_adr[hid] - normal_base = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), quat) + mat_col = rot_vec_quat(wp.vec3(0.0, 0.0, 1.0), quat) # compute size and pos of base box base_scale = size[3] * 0.5 base_size = wp.vec3(size[0], size[1], base_scale) - base_pos = pos + normal_base * base_scale + base_pos = pos + mat_col * base_scale # compute size and pos of top box top_scale = size[2] * 0.5 top_size = wp.vec3(size[0], size[1], top_scale) - top_pos = pos + normal_base * top_scale + top_pos = pos + mat_col * top_scale # init: intersection with base box - x, _ = _ray_box(base_pos, quat, base_size, pnt, vec) + x, _, normal_base = _ray_box(base_pos, quat, base_size, pnt, vec) # check top box: done if no intersection top_intersect, all, normal_top = _ray_box(top_pos, quat, top_size, pnt, vec) From 75235a71bfa419286a093d90745dad3855fd1674 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 7 Jan 2026 08:37:09 +0100 Subject: [PATCH 103/107] whitespace fix Signed-off-by: Alain Denzler --- mujoco_warp/_src/types.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 1a65d6d4c..4566cdf01 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -1643,9 +1643,9 @@ class Data: collision_pairid: ids from broadphase (naconmax, 2) collision_worldid: collision world ids from broadphase (naconmax,) ncollision: collision count from broadphase (1,) - xiquat: Cartesian orientation of body inertia (nworld, nbody, 4) - geom_xquat: Cartesian orientation of geom (nworld, ngeom, 4) - site_xquat: Cartesian orientation of site (nworld, nsite, 4) + xiquat: Cartesian orientation of body inertia (nworld, nbody, 4) + geom_xquat: Cartesian orientation of geom (nworld, ngeom, 4) + site_xquat: Cartesian orientation of site (nworld, nsite, 4) """ solver_niter: array("nworld", int) From 739eca0c948e4833219e8b025a692829e842f6e3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 7 Jan 2026 08:44:21 +0100 Subject: [PATCH 104/107] missing mat->quat conversions Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 2 +- mujoco_warp/_src/sensor.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index e338aefda..7fc17b059 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -106,7 +106,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint: if geomtype == GeomType.BOX: tmp = wp.sign(local_dir) res = wp.cw_mul(tmp, geom.size) - sp.point = quat_to_mat(geom.rot) @ res + geom.pos + sp.point = rot_vec_quat(res, geom.rot) + geom.pos sp.vertex_index = wp.where(tmp[0] > 0, 1, 0) sp.vertex_index += wp.where(tmp[1] > 0, 2, 0) sp.vertex_index += wp.where(tmp[2] > 0, 4, 0) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 26626cc5f..4237c993a 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -2175,7 +2175,7 @@ def _sensor_tactile( # position in other geom frame tmp = xpos - geom_xpos_in[worldid, geom] - lpos = wp.transpose(math.quat_to_mat(quat)) @ tmp + lpos = math.rot_vec_quat(tmp, math.quat_inv(quat)) plugin_id = geom_plugin_index[geom] From 62ea9ac8c45a44370150e32e2a2327ffce1f109c Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 7 Jan 2026 08:46:23 +0100 Subject: [PATCH 105/107] import cleanup Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 1 - 1 file changed, 1 deletion(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index f157a8a62..b4ea7ba2a 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -19,7 +19,6 @@ from .collision_primitive import Geom from .math import quat_inv -from .math import quat_to_mat from .math import rot_vec_quat from .types import GeomType from .types import mat43 From a9dd5cf4e52ec130f8353eeba414075ced04bdf1 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 Jan 2026 09:22:41 +0100 Subject: [PATCH 106/107] change epsilon in half-space calc Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_gjk.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mujoco_warp/_src/collision_gjk.py b/mujoco_warp/_src/collision_gjk.py index b4ea7ba2a..7b491016e 100644 --- a/mujoco_warp/_src/collision_gjk.py +++ b/mujoco_warp/_src/collision_gjk.py @@ -1807,7 +1807,7 @@ def _plane_normal(v1: wp.vec3, v2: wp.vec3, n: wp.vec3) -> Tuple[float, wp.vec3] @wp.func def _halfspace(a: wp.vec3, n: wp.vec3, p: wp.vec3) -> bool: - return wp.dot(p - a, n) > -1e-10 + return wp.dot(p - a, n) > -1e-6 @wp.func From c674bf760b9607946f19b3a440a5c10f9cc82f8b Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 12 Jan 2026 09:58:14 +0100 Subject: [PATCH 107/107] heightfield merge conflicts Signed-off-by: Alain Denzler --- mujoco_warp/_src/collision_convex.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 1d62b0dc7..e6c14ad4d 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -326,15 +326,15 @@ def ccd_hfield_kernel( # transform geom2 into heightfield frame hf_pos = geom_xpos_in[worldid, g1] - hf_mat = geom_xmat_in[worldid, g1] - hf_matT = wp.transpose(hf_mat) + hf_quat = geom_xquat_in[worldid, g1] + hf_quat_inv = quat_inv(hf_quat) - geom2.pos = hf_matT @ (geom2.pos - hf_pos) - geom2.rot = hf_matT @ geom2.rot + geom2.pos = rot_vec_quat(geom2.pos - hf_pos, hf_quat_inv) + geom2.rot = mul_quat(hf_quat_inv, geom2.rot) # geom1 has identity pose geom1.pos = wp.vec3(0.0, 0.0, 0.0) - geom1.rot = wp.identity(n=3, dtype=float) + geom1.rot = wp.quat(1.0, 0.0, 0.0, 0.0) # see MuJoCo mjc_ConvexHField geom1_dataid = geom_dataid[g1] @@ -479,14 +479,14 @@ def ccd_hfield_kernel( # transform contact to global frame pos_local = 0.5 * (w1 + w2) - pos = hf_mat @ pos_local + hf_pos + pos = rot_vec_quat(pos_local, hf_quat) + hf_pos hfield_contact_pos[count, 0] = pos[0] hfield_contact_pos[count, 1] = pos[1] hfield_contact_pos[count, 2] = pos[2] frame_local = make_frame(w1 - w2) normal_local = wp.vec3(frame_local[0, 0], frame_local[0, 1], frame_local[0, 2]) - normal = hf_mat @ normal_local + normal = rot_vec_quat(normal_local, hf_quat) hfield_contact_normal[count, 0] = normal[0] hfield_contact_normal[count, 1] = normal[1] hfield_contact_normal[count, 2] = normal[2]