From 02f3485a2ea4f75d019756a84e45ad254f594f58 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 16:35:36 +0200 Subject: [PATCH 01/74] io.py changes. just a proposal. --- mujoco_warp/_src/io.py | 147 +++++++++++++++++++++-------------------- 1 file changed, 77 insertions(+), 70 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index 513cc1f0..cf5b817d 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -246,6 +246,13 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: nxn_geom_pair.append((geom1, geom2)) nxn_pairid.append(pairid) + def create_nmodel_batched_array(mjm_array, dtype): + array = wp.array(mjm_array, dtype=dtype) + array.ndim += 1 + array.shape = (1,) + array.shape + array.strides = (0,) + array.strides + return array + return types.Model( nq=mjm.nq, nv=mjm.nv, @@ -290,8 +297,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: stat=types.Statistic( meaninertia=mjm.stat.meaninertia, ), - qpos0=wp.array(mjm.qpos0, dtype=float), - qpos_spring=wp.array(mjm.qpos_spring, dtype=float), + qpos0=create_nmodel_batched_array(mjm.qpos0, dtype=float), + qpos_spring=create_nmodel_batched_array(mjm.qpos_spring, dtype=float), qM_fullm_i=wp.array(qM_fullm_i, dtype=int), qM_fullm_j=wp.array(qM_fullm_j, dtype=int), qM_mulm_i=wp.array(qM_mulm_i, dtype=int), @@ -314,15 +321,15 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: body_dofadr=wp.array(mjm.body_dofadr, dtype=int), body_geomnum=wp.array(mjm.body_geomnum, dtype=int), body_geomadr=wp.array(mjm.body_geomadr, dtype=int), - body_pos=wp.array(mjm.body_pos, dtype=wp.vec3), - body_quat=wp.array(mjm.body_quat, dtype=wp.quat), - body_ipos=wp.array(mjm.body_ipos, dtype=wp.vec3), - body_iquat=wp.array(mjm.body_iquat, dtype=wp.quat), - body_mass=wp.array(mjm.body_mass, dtype=float), - body_subtreemass=wp.array(mjm.body_subtreemass, dtype=float), - subtree_mass=wp.array(subtree_mass, dtype=float), - body_inertia=wp.array(mjm.body_inertia, dtype=wp.vec3), - body_invweight0=wp.array(mjm.body_invweight0, dtype=float), + body_pos=create_nmodel_batched_array(mjm.body_pos, dtype=wp.vec3), + body_quat=create_nmodel_batched_array(mjm.body_quat, dtype=wp.quat), + body_ipos=create_nmodel_batched_array(mjm.body_ipos, dtype=wp.vec3), + body_iquat=create_nmodel_batched_array(mjm.body_iquat, dtype=wp.quat), + body_mass=create_nmodel_batched_array(mjm.body_mass, dtype=float), + body_subtreemass=create_nmodel_batched_array(mjm.body_subtreemass, dtype=float), + subtree_mass=create_nmodel_batched_array(subtree_mass, dtype=float), + body_inertia=create_nmodel_batched_array(mjm.body_inertia, dtype=wp.vec3), + body_invweight0=create_nmodel_batched_array(mjm.body_invweight0, dtype=float), body_contype=wp.array(mjm.body_contype, dtype=int), body_conaffinity=wp.array(mjm.body_conaffinity, dtype=int), body_gravcomp=wp.array(mjm.body_gravcomp, dtype=float), @@ -331,15 +338,15 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: jnt_dofadr=wp.array(mjm.jnt_dofadr, dtype=int), jnt_bodyid=wp.array(mjm.jnt_bodyid, dtype=int), jnt_limited=wp.array(mjm.jnt_limited, dtype=int), - jnt_actfrclimited=wp.array(mjm.jnt_actfrclimited, dtype=bool), - jnt_solref=wp.array(mjm.jnt_solref, dtype=wp.vec2), - jnt_solimp=wp.array(mjm.jnt_solimp, dtype=types.vec5), - jnt_pos=wp.array(mjm.jnt_pos, dtype=wp.vec3), + jnt_actfrclimited=create_nmodel_batched_array(mjm.jnt_actfrclimited, dtype=bool), + jnt_solref=create_nmodel_batched_array(mjm.jnt_solref, dtype=wp.vec2), + jnt_solimp=create_nmodel_batched_array(mjm.jnt_solimp, dtype=types.vec5), + jnt_pos=create_nmodel_batched_array(mjm.jnt_pos, dtype=wp.vec3), jnt_axis=wp.array(mjm.jnt_axis, dtype=wp.vec3), - jnt_stiffness=wp.array(mjm.jnt_stiffness, dtype=float), - jnt_range=wp.array(mjm.jnt_range, dtype=float), - jnt_actfrcrange=wp.array(mjm.jnt_actfrcrange, dtype=wp.vec2), - jnt_margin=wp.array(mjm.jnt_margin, dtype=float), + jnt_stiffness=create_nmodel_batched_array(mjm.jnt_stiffness, dtype=float), + jnt_range=create_nmodel_batched_array(mjm.jnt_range, dtype=float), + jnt_actfrcrange=create_nmodel_batched_array(mjm.jnt_actfrcrange, dtype=wp.vec2), + jnt_margin=create_nmodel_batched_array(mjm.jnt_margin, dtype=float), # these jnt_limited adrs are used in constraint.py jnt_limited_slide_hinge_adr=wp.array( np.nonzero( @@ -356,12 +363,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: dof_jntid=wp.array(mjm.dof_jntid, dtype=int), dof_parentid=wp.array(mjm.dof_parentid, dtype=int), dof_Madr=wp.array(mjm.dof_Madr, dtype=int), - dof_armature=wp.array(mjm.dof_armature, dtype=float), - dof_damping=wp.array(mjm.dof_damping, dtype=float), - dof_invweight0=wp.array(mjm.dof_invweight0, dtype=float), - dof_frictionloss=wp.array(mjm.dof_frictionloss, dtype=float), - dof_solimp=wp.array(mjm.dof_solimp, dtype=types.vec5), - dof_solref=wp.array(mjm.dof_solref, dtype=wp.vec2), + dof_armature=create_nmodel_batched_array(mjm.dof_armature, dtype=float), + dof_damping=create_nmodel_batched_array(mjm.dof_damping, dtype=float), + dof_invweight0=create_nmodel_batched_array(mjm.dof_invweight0, dtype=float), + dof_frictionloss=create_nmodel_batched_array(mjm.dof_frictionloss, dtype=float), + dof_solimp=create_nmodel_batched_array(mjm.dof_solimp, dtype=types.vec5), + dof_solref=create_nmodel_batched_array(mjm.dof_solref, dtype=wp.vec2), dof_tri_row=wp.array(dof_tri_row, dtype=int), dof_tri_col=wp.array(dof_tri_col, dtype=int), geom_type=wp.array(mjm.geom_type, dtype=int), @@ -370,28 +377,28 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: geom_condim=wp.array(mjm.geom_condim, dtype=int), geom_bodyid=wp.array(mjm.geom_bodyid, dtype=int), geom_dataid=wp.array(mjm.geom_dataid, dtype=int), - geom_priority=wp.array(mjm.geom_priority, dtype=int), - geom_solmix=wp.array(mjm.geom_solmix, dtype=float), - geom_solref=wp.array(mjm.geom_solref, dtype=wp.vec2), - geom_solimp=wp.array(mjm.geom_solimp, dtype=types.vec5), - geom_size=wp.array(mjm.geom_size, dtype=wp.vec3), - geom_aabb=wp.array(mjm.geom_aabb, dtype=wp.vec3), - geom_rbound=wp.array(mjm.geom_rbound, dtype=float), - geom_pos=wp.array(mjm.geom_pos, dtype=wp.vec3), - geom_quat=wp.array(mjm.geom_quat, dtype=wp.quat), - geom_friction=wp.array(mjm.geom_friction, dtype=wp.vec3), - geom_margin=wp.array(mjm.geom_margin, dtype=float), - geom_gap=wp.array(mjm.geom_gap, dtype=float), + geom_priority=create_nmodel_batched_array(mjm.geom_priority, dtype=int), + geom_solmix=create_nmodel_batched_array(mjm.geom_solmix, dtype=float), + geom_solref=create_nmodel_batched_array(mjm.geom_solref, dtype=wp.vec2), + geom_solimp=create_nmodel_batched_array(mjm.geom_solimp, dtype=types.vec5), + geom_size=create_nmodel_batched_array(mjm.geom_size, dtype=wp.vec3), + geom_aabb=create_nmodel_batched_array(mjm.geom_aabb, dtype=wp.vec3), + geom_rbound=create_nmodel_batched_array(mjm.geom_rbound, dtype=float), + geom_pos=create_nmodel_batched_array(mjm.geom_pos, dtype=wp.vec3), + geom_quat=create_nmodel_batched_array(mjm.geom_quat, dtype=wp.quat), + geom_friction=create_nmodel_batched_array(mjm.geom_friction, dtype=wp.vec3), + geom_margin=create_nmodel_batched_array(mjm.geom_margin, dtype=float), + geom_gap=create_nmodel_batched_array(mjm.geom_gap, dtype=float), site_bodyid=wp.array(mjm.site_bodyid, dtype=int), - site_pos=wp.array(mjm.site_pos, dtype=wp.vec3), - site_quat=wp.array(mjm.site_quat, dtype=wp.quat), + site_pos=create_nmodel_batched_array(mjm.site_pos, dtype=wp.vec3), + site_quat=create_nmodel_batched_array(mjm.site_quat, dtype=wp.quat), cam_mode=wp.array(mjm.cam_mode, dtype=int), cam_bodyid=wp.array(mjm.cam_bodyid, dtype=int), cam_targetbodyid=wp.array(mjm.cam_targetbodyid, dtype=int), - cam_pos=wp.array(mjm.cam_pos, dtype=wp.vec3), - cam_quat=wp.array(mjm.cam_quat, dtype=wp.quat), - cam_poscom0=wp.array(mjm.cam_poscom0, dtype=wp.vec3), - cam_pos0=wp.array(mjm.cam_pos0, dtype=wp.vec3), + cam_pos=create_nmodel_batched_array(mjm.cam_pos, dtype=wp.vec3), + cam_quat=create_nmodel_batched_array(mjm.cam_quat, dtype=wp.quat), + cam_poscom0=create_nmodel_batched_array(mjm.cam_poscom0, dtype=wp.vec3), + cam_pos0=create_nmodel_batched_array(mjm.cam_pos0, dtype=wp.vec3), cam_fovy=wp.array(mjm.cam_fovy, dtype=float), cam_resolution=wp.array(mjm.cam_resolution, dtype=wp.vec2i), cam_sensorsize=wp.array(mjm.cam_sensorsize, dtype=wp.vec2), @@ -399,10 +406,10 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: light_mode=wp.array(mjm.light_mode, dtype=int), light_bodyid=wp.array(mjm.light_bodyid, dtype=int), light_targetbodyid=wp.array(mjm.light_targetbodyid, dtype=int), - light_pos=wp.array(mjm.light_pos, dtype=wp.vec3), - light_dir=wp.array(mjm.light_dir, dtype=wp.vec3), - light_poscom0=wp.array(mjm.light_poscom0, dtype=wp.vec3), - light_pos0=wp.array(mjm.light_pos0, dtype=wp.vec3), + light_pos=create_nmodel_batched_array(mjm.light_pos, dtype=wp.vec3), + light_dir=create_nmodel_batched_array(mjm.light_dir, dtype=wp.vec3), + light_poscom0=create_nmodel_batched_array(mjm.light_poscom0, dtype=wp.vec3), + light_pos0=create_nmodel_batched_array(mjm.light_pos0, dtype=wp.vec3), mesh_vertadr=wp.array(mjm.mesh_vertadr, dtype=int), mesh_vertnum=wp.array(mjm.mesh_vertnum, dtype=int), mesh_vert=wp.array(mjm.mesh_vert, dtype=wp.vec3), @@ -411,8 +418,8 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: eq_obj2id=wp.array(mjm.eq_obj2id, dtype=int), eq_objtype=wp.array(mjm.eq_objtype, dtype=int), eq_active0=wp.array(mjm.eq_active0, dtype=bool), - eq_solref=wp.array(mjm.eq_solref, dtype=wp.vec2), - eq_solimp=wp.array(mjm.eq_solimp, dtype=types.vec5), + eq_solref=create_nmodel_batched_array(mjm.eq_solref, dtype=wp.vec2), + eq_solimp=create_nmodel_batched_array(mjm.eq_solimp, dtype=types.vec5), eq_data=wp.array(mjm.eq_data, dtype=types.vec11), # pre-compute indices of equality constraints eq_connect_adr=wp.array(np.nonzero(mjm.eq_type == types.EqType.CONNECT.value)[0], dtype=int), @@ -431,13 +438,13 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: actuator_ctrllimited=wp.array(mjm.actuator_ctrllimited, dtype=bool), actuator_forcelimited=wp.array(mjm.actuator_forcelimited, dtype=bool), actuator_actlimited=wp.array(mjm.actuator_actlimited, dtype=bool), - actuator_dynprm=wp.array(mjm.actuator_dynprm, dtype=types.vec10f), - actuator_gainprm=wp.array(mjm.actuator_gainprm, dtype=types.vec10f), - actuator_biasprm=wp.array(mjm.actuator_biasprm, dtype=types.vec10f), - actuator_ctrlrange=wp.array(mjm.actuator_ctrlrange, dtype=wp.vec2), - actuator_forcerange=wp.array(mjm.actuator_forcerange, dtype=wp.vec2), - actuator_actrange=wp.array(mjm.actuator_actrange, dtype=wp.vec2), - actuator_gear=wp.array(mjm.actuator_gear, dtype=wp.spatial_vector), + actuator_dynprm=create_nmodel_batched_array(mjm.actuator_dynprm, dtype=types.vec10f), + actuator_gainprm=create_nmodel_batched_array(mjm.actuator_gainprm, dtype=types.vec10f), + actuator_biasprm=create_nmodel_batched_array(mjm.actuator_biasprm, dtype=types.vec10f), + actuator_ctrlrange=create_nmodel_batched_array(mjm.actuator_ctrlrange, dtype=wp.vec2), + actuator_forcerange=create_nmodel_batched_array(mjm.actuator_forcerange, dtype=wp.vec2), + actuator_actrange=create_nmodel_batched_array(mjm.actuator_actrange, dtype=wp.vec2), + actuator_gear=create_nmodel_batched_array(mjm.actuator_gear, dtype=wp.spatial_vector), exclude_signature=wp.array(mjm.exclude_signature, dtype=int), # short-circuiting here allows us to skip a lot of code in implicit integration actuator_affine_bias_gain=bool( @@ -449,25 +456,25 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: pair_dim=wp.array(mjm.pair_dim, dtype=int), pair_geom1=wp.array(mjm.pair_geom1, dtype=int), pair_geom2=wp.array(mjm.pair_geom2, dtype=int), - pair_solref=wp.array(mjm.pair_solref, dtype=wp.vec2), - pair_solreffriction=wp.array(mjm.pair_solreffriction, dtype=wp.vec2), - pair_solimp=wp.array(mjm.pair_solimp, dtype=types.vec5), - pair_margin=wp.array(mjm.pair_margin, dtype=float), - pair_gap=wp.array(mjm.pair_gap, dtype=float), - pair_friction=wp.array(mjm.pair_friction, dtype=types.vec5), + pair_solref=create_nmodel_batched_array(mjm.pair_solref, dtype=wp.vec2), + pair_solreffriction=create_nmodel_batched_array(mjm.pair_solreffriction, dtype=wp.vec2), + pair_solimp=create_nmodel_batched_array(mjm.pair_solimp, dtype=types.vec5), + pair_margin=create_nmodel_batched_array(mjm.pair_margin, dtype=float), + pair_gap=create_nmodel_batched_array(mjm.pair_gap, dtype=float), + pair_friction=create_nmodel_batched_array(mjm.pair_friction, dtype=types.vec5), condim_max=np.max(mjm.pair_dim) if mjm.npair else np.max(mjm.geom_condim), # TODO(team): get max after filtering, tendon_adr=wp.array(mjm.tendon_adr, dtype=int), tendon_num=wp.array(mjm.tendon_num, dtype=int), tendon_limited=wp.array(mjm.tendon_limited, dtype=int), tendon_limited_adr=wp.array(np.nonzero(mjm.tendon_limited)[0], dtype=wp.int32, ndim=1), - tendon_solref_lim=wp.array(mjm.tendon_solref_lim, dtype=wp.vec2f), - tendon_solimp_lim=wp.array(mjm.tendon_solimp_lim, dtype=types.vec5), - tendon_range=wp.array(mjm.tendon_range, dtype=wp.vec2f), - tendon_margin=wp.array(mjm.tendon_margin, dtype=float), - tendon_length0=wp.array(mjm.tendon_length0, dtype=float), - tendon_invweight0=wp.array(mjm.tendon_invweight0, dtype=float), + tendon_solref_lim=create_nmodel_batched_array(mjm.tendon_solref_lim, dtype=wp.vec2f), + tendon_solimp_lim=create_nmodel_batched_array(mjm.tendon_solimp_lim, dtype=types.vec5), + tendon_range=create_nmodel_batched_array(mjm.tendon_range, dtype=wp.vec2f), + tendon_margin=create_nmodel_batched_array(mjm.tendon_margin, dtype=float), + tendon_length0=create_nmodel_batched_array(mjm.tendon_length0, dtype=float), + tendon_invweight0=create_nmodel_batched_array(mjm.tendon_invweight0, dtype=float), wrap_objid=wp.array(mjm.wrap_objid, dtype=int), - wrap_prm=wp.array(mjm.wrap_prm, dtype=float), + wrap_prm=create_nmodel_batched_array(mjm.wrap_prm, dtype=float), wrap_type=wp.array(mjm.wrap_type, dtype=int), tendon_jnt_adr=wp.array(tendon_jnt_adr, dtype=int), tendon_site_adr=wp.array(tendon_site_adr, dtype=int), @@ -485,7 +492,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model: sensor_refid=wp.array(mjm.sensor_refid, dtype=int), sensor_dim=wp.array(mjm.sensor_dim, dtype=int), sensor_adr=wp.array(mjm.sensor_adr, dtype=int), - sensor_cutoff=wp.array(mjm.sensor_cutoff, dtype=float), + sensor_cutoff=create_nmodel_batched_array(mjm.sensor_cutoff, dtype=float), sensor_pos_adr=wp.array( np.nonzero(mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS)[0], dtype=int, From 569995df3d49b1e70ee2e43f465cf8d273de3771 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 16:46:39 +0200 Subject: [PATCH 02/74] docs in types.py --- mujoco_warp/_src/io.py | 2 +- mujoco_warp/_src/types.py | 136 +++++++++++++++++++------------------- 2 files changed, 69 insertions(+), 69 deletions(-) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index cf5b817d..946c66b8 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -338,7 +338,7 @@ def create_nmodel_batched_array(mjm_array, dtype): jnt_dofadr=wp.array(mjm.jnt_dofadr, dtype=int), jnt_bodyid=wp.array(mjm.jnt_bodyid, dtype=int), jnt_limited=wp.array(mjm.jnt_limited, dtype=int), - jnt_actfrclimited=create_nmodel_batched_array(mjm.jnt_actfrclimited, dtype=bool), + jnt_actfrclimited=wp.array(mjm.jnt_actfrclimited, dtype=bool), jnt_solref=create_nmodel_batched_array(mjm.jnt_solref, dtype=wp.vec2), jnt_solimp=create_nmodel_batched_array(mjm.jnt_solimp, dtype=types.vec5), jnt_pos=create_nmodel_batched_array(mjm.jnt_pos, dtype=wp.vec3), diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index f87ea93e..58aebd0f 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -587,8 +587,8 @@ class Model: nlsp: number of step sizes for parallel linsearch () opt: physics options stat: model statistics - qpos0: qpos values at default pose (nq,) - qpos_spring: reference pose for springs (nq,) + qpos0: qpos values at default pose (nworld, nq) + qpos_spring: reference pose for springs (nworld, nq) qM_fullm_i: sparse mass matrix addressing qM_fullm_j: sparse mass matrix addressing qM_mulm_i: sparse mass matrix addressing @@ -612,15 +612,15 @@ class Model: body_dofadr: start addr of dofs; -1: no dofs (nbody,) body_geomnum: number of geoms (nbody,) body_geomadr: start addr of geoms; -1: no geoms (nbody,) - body_pos: position offset rel. to parent body (nbody, 3) - body_quat: orientation offset rel. to parent body (nbody, 4) - body_ipos: local position of center of mass (nbody, 3) - body_iquat: local orientation of inertia ellipsoid (nbody, 4) - body_mass: mass (nbody,) - body_subtreemass: mass of subtree starting at this body (nbody,) - subtree_mass: mass of subtree (nbody,) - body_inertia: diagonal inertia in ipos/iquat frame (nbody, 3) - body_invweight0: mean inv inert in qpos0 (trn, rot) (nbody, 2) + body_pos: position offset rel. to parent body (nworld, nbody, 3) + body_quat: orientation offset rel. to parent body (nworld, nbody, 4) + body_ipos: local position of center of mass (nworld, nbody, 3) + body_iquat: local orientation of inertia ellipsoid (nworld, nbody, 4) + body_mass: mass (nworld, nbody,) + body_subtreemass: mass of subtree starting at this body (nworld, nbody,) + subtree_mass: mass of subtree (nworld, nbody,) + body_inertia: diagonal inertia in ipos/iquat frame (nworld, nbody, 3) + body_invweight0: mean inv inert in qpos0 (trn, rot) (nworld, nbody, 2) body_contype: OR over all geom contypes (nbody,) body_conaffinity: OR over all geom conaffinities (nbody,) body_gravcomp: antigravity force, units of body weight (nbody,) @@ -630,14 +630,14 @@ class Model: jnt_bodyid: id of joint's body (njnt,) jnt_limited: does joint have limits (njnt,) jnt_actfrclimited: does joint have actuator force limits (njnt,) - jnt_solref: constraint solver reference: limit (njnt, mjNREF) - jnt_solimp: constraint solver impedance: limit (njnt, mjNIMP) - jnt_pos: local anchor position (njnt, 3) + jnt_solref: constraint solver reference: limit (nworld, njnt, mjNREF) + jnt_solimp: constraint solver impedance: limit (nworld, njnt, mjNIMP) + jnt_pos: local anchor position (nworld,njnt, 3) jnt_axis: local joint axis (njnt, 3) - jnt_stiffness: stiffness coefficient (njnt,) - jnt_range: joint limits (njnt, 2) - jnt_actfrcrange: range of total actuator force (njnt, 2) - jnt_margin: min distance for limit detection (njnt,) + jnt_stiffness: stiffness coefficient (nworld, njnt) + jnt_range: joint limits (nworld, njnt, 2) + jnt_actfrcrange: range of total actuator force (nworldnjnt, 2) + jnt_margin: min distance for limit detection (nworld, njnt) jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr jnt_limited_ball_adr: limited/ball jntadr jnt_actgravcomp: is gravcomp force applied via actuators (njnt,) @@ -645,12 +645,12 @@ class Model: dof_jntid: id of dof's joint (nv,) dof_parentid: id of dof's parent; -1: none (nv,) dof_Madr: dof address in M-diagonal (nv,) - dof_armature: dof armature inertia/mass (nv,) - dof_damping: damping coefficient (nv,) - dof_invweight0: diag. inverse inertia in qpos0 (nv,) - dof_frictionloss: dof friction loss (nv,) - dof_solimp: constraint solver impedance: frictionloss (nv, NIMP) - dof_solref: constraint solver reference: frictionloss (nv, NREF) + dof_armature: dof armature inertia/mass (nworld, nv) + dof_damping: damping coefficient (nworld, nv) + dof_invweight0: diag. inverse inertia in qpos0 (nworld, nv) + dof_frictionloss: dof friction loss (nworld, nv) + dof_solimp: constraint solver impedance: frictionloss (nworld, nv, NIMP) + dof_solref: constraint solver reference: frictionloss (nworld,nv, NREF) dof_tri_row: np.tril_indices (mjm.nv)[0] dof_tri_col: np.tril_indices (mjm.nv)[1] geom_type: geometric type (mjtGeom) (ngeom,) @@ -659,28 +659,28 @@ class Model: geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,) geom_bodyid: id of geom's body (ngeom,) geom_dataid: id of geom's mesh/hfield; -1: none (ngeom,) - geom_priority: geom contact priority (ngeom,) - geom_solmix: mixing coef for solref/imp in geom pair (ngeom,) - geom_solref: constraint solver reference: contact (ngeom, mjNREF) - geom_solimp: constraint solver impedance: contact (ngeom, mjNIMP) - geom_size: geom-specific size parameters (ngeom, 3) - geom_aabb: bounding box, (center, size) (ngeom, 6) - geom_rbound: radius of bounding sphere (ngeom,) - geom_pos: local position offset rel. to body (ngeom, 3) - geom_quat: local orientation offset rel. to body (ngeom, 4) - geom_friction: friction for (slide, spin, roll) (ngeom, 3) - geom_margin: detect contact if dist Date: Thu, 8 May 2025 16:52:01 +0200 Subject: [PATCH 03/74] qpos0 --- mujoco_warp/_src/constraint.py | 8 ++++---- mujoco_warp/_src/smooth.py | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index f7b5481b..f193dfe7 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -231,7 +231,7 @@ def _efc_equality_connect( def _efc_equality_joint( # Model: opt_timestep: float, - qpos0: wp.array(dtype=float), + qpos0: wp.array2d(dtype=float), jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), dof_invweight0: wp.array(dtype=float), @@ -281,20 +281,20 @@ def _efc_equality_joint( # Two joint constraint qposadr2 = jnt_qposadr[jntid_2] dofadr2 = jnt_dofadr[jntid_2] - dif = qpos_in[worldid, qposadr2] - qpos0[qposadr2] + dif = qpos_in[worldid, qposadr2] - qpos0[worldid, qposadr2] # Horner's method for polynomials rhs = data[0] + dif * (data[1] + dif * (data[2] + dif * (data[3] + dif * data[4]))) deriv_2 = data[1] + dif * (2.0 * data[2] + dif * (3.0 * data[3] + dif * 4.0 * data[4])) - pos = qpos_in[worldid, qposadr1] - qpos0[qposadr1] - rhs + pos = qpos_in[worldid, qposadr1] - qpos0[worldid,qposadr1] - rhs Jqvel = qvel_in[worldid, dofadr1] - qvel_in[worldid, dofadr2] * deriv_2 invweight = dof_invweight0[dofadr1] + dof_invweight0[dofadr2] efc_J_out[efcid, dofadr2] = -deriv_2 else: # Single joint constraint - pos = qpos_in[worldid, qposadr1] - qpos0[qposadr1] - data[0] + pos = qpos_in[worldid, qposadr1] - qpos0[worldid, qposadr1] - data[0] Jqvel = qvel_in[worldid, dofadr1] invweight = dof_invweight0[dofadr1] diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 0a7ee296..72b05cf3 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -60,7 +60,7 @@ def _kinematics_root( @wp.kernel def _kinematics_level( # Model: - qpos0: wp.array(dtype=float), + qpos0: wp.array2d(dtype=float), body_parentid: wp.array(dtype=int), body_jntnum: wp.array(dtype=int), body_jntadr: wp.array(dtype=int), @@ -131,9 +131,9 @@ def _kinematics_level( # correct for off-center rotation xpos = xanchor - math.rot_vec_quat(jnt_pos[jntadr], xquat) elif jnt_type_ == wp.static(JointType.SLIDE.value): - xpos += xaxis * (qpos[qadr] - qpos0[qadr]) + xpos += xaxis * (qpos[qadr] - qpos0[worldid, qadr]) elif jnt_type_ == wp.static(JointType.HINGE.value): - qpos0_ = qpos0[qadr] + qpos0_ = qpos0[worldid, qadr] qloc_ = math.axis_angle_to_quat(jnt_axis_, qpos[qadr] - qpos0_) xquat = math.mul_quat(xquat, qloc_) # correct for off-center rotation From 1c3930dd50065edf19b39c1414122afcb806c8ac Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 16:53:08 +0200 Subject: [PATCH 04/74] qpos_spring --- mujoco_warp/_src/passive.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index c382716d..1f82ec13 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -27,7 +27,7 @@ @wp.kernel def _spring_passive( # Model: - qpos_spring: wp.array(dtype=float), + qpos_spring: wp.array2d(dtype=float), jnt_type: wp.array(dtype=int), jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), @@ -49,9 +49,9 @@ def _spring_passive( if jnttype == wp.static(JointType.FREE.value): dif = wp.vec3( - qpos_in[worldid, qposid + 0] - qpos_spring[qposid + 0], - qpos_in[worldid, qposid + 1] - qpos_spring[qposid + 1], - qpos_in[worldid, qposid + 2] - qpos_spring[qposid + 2], + qpos_in[worldid, qposid + 0] - qpos_spring[worldid, qposid + 0], + qpos_in[worldid, qposid + 1] - qpos_spring[worldid, qposid + 1], + qpos_in[worldid, qposid + 2] - qpos_spring[worldid, qposid + 2], ) qfrc_spring_out[worldid, dofid + 0] = -stiffness * dif[0] qfrc_spring_out[worldid, dofid + 1] = -stiffness * dif[1] @@ -63,10 +63,10 @@ def _spring_passive( qpos_in[worldid, qposid + 6], ) ref = wp.quat( - qpos_spring[qposid + 3], - qpos_spring[qposid + 4], - qpos_spring[qposid + 5], - qpos_spring[qposid + 6], + qpos_spring[worldid, qposid + 3], + qpos_spring[worldid, qposid + 4], + qpos_spring[worldid, qposid + 5], + qpos_spring[worldid, qposid + 6], ) dif = math.quat_sub(rot, ref) qfrc_spring_out[worldid, dofid + 3] = -stiffness * dif[0] @@ -80,17 +80,17 @@ def _spring_passive( qpos_in[worldid, qposid + 3], ) ref = wp.quat( - qpos_spring[qposid + 0], - qpos_spring[qposid + 1], - qpos_spring[qposid + 2], - qpos_spring[qposid + 3], + qpos_spring[worldid, qposid + 0], + qpos_spring[worldid, qposid + 1], + qpos_spring[worldid, qposid + 2], + qpos_spring[worldid, qposid + 3], ) dif = math.quat_sub(rot, ref) qfrc_spring_out[worldid, dofid + 0] = -stiffness * dif[0] qfrc_spring_out[worldid, dofid + 1] = -stiffness * dif[1] qfrc_spring_out[worldid, dofid + 2] = -stiffness * dif[2] else: # mjJNT_SLIDE, mjJNT_HINGE - fdif = qpos_in[worldid, qposid] - qpos_spring[qposid] + fdif = qpos_in[worldid, qposid] - qpos_spring[worldid,qposid] qfrc_spring_out[worldid, dofid] = -stiffness * fdif From f6bea97b9a34674844d457edb63d75b38a69f62b Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 16:54:39 +0200 Subject: [PATCH 05/74] body_pos --- 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 72b05cf3..b43a3bde 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -64,7 +64,7 @@ def _kinematics_level( body_parentid: wp.array(dtype=int), body_jntnum: wp.array(dtype=int), body_jntadr: wp.array(dtype=int), - body_pos: wp.array(dtype=wp.vec3), + body_pos: wp.array2d(dtype=wp.vec3), body_quat: wp.array(dtype=wp.quat), body_ipos: wp.array(dtype=wp.vec3), body_iquat: wp.array(dtype=wp.quat), @@ -97,7 +97,7 @@ def _kinematics_level( if jntnum == 0: # no joints - apply fixed translation and rotation relative to parent pid = body_parentid[bodyid] - xpos = (xmat_in[worldid, pid] * body_pos[bodyid]) + xpos_in[worldid, pid] + xpos = (xmat_in[worldid, pid] * body_pos[worldid, bodyid]) + xpos_in[worldid, pid] xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[bodyid]) elif jntnum == 1 and jnt_type[jntadr] == wp.static(JointType.FREE.value): # free joint @@ -110,7 +110,7 @@ def _kinematics_level( # regular or no joints # apply fixed translation and rotation relative to parent pid = body_parentid[bodyid] - xpos = (xmat_in[worldid, pid] * body_pos[bodyid]) + xpos_in[worldid, pid] + xpos = (xmat_in[worldid, pid] * body_pos[worldid, bodyid]) + xpos_in[worldid, pid] xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[bodyid]) for _ in range(jntnum): From f8c62d014124d9269250a1395be938978a2571a3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 16:55:12 +0200 Subject: [PATCH 06/74] body_quat --- 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 b43a3bde..f088a1a7 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -65,7 +65,7 @@ def _kinematics_level( body_jntnum: wp.array(dtype=int), body_jntadr: wp.array(dtype=int), body_pos: wp.array2d(dtype=wp.vec3), - body_quat: wp.array(dtype=wp.quat), + body_quat: wp.array2d(dtype=wp.quat), body_ipos: wp.array(dtype=wp.vec3), body_iquat: wp.array(dtype=wp.quat), jnt_type: wp.array(dtype=int), @@ -98,7 +98,7 @@ def _kinematics_level( # no joints - apply fixed translation and rotation relative to parent pid = body_parentid[bodyid] xpos = (xmat_in[worldid, pid] * body_pos[worldid, bodyid]) + xpos_in[worldid, pid] - xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[bodyid]) + xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[worldid, bodyid]) elif jntnum == 1 and jnt_type[jntadr] == wp.static(JointType.FREE.value): # free joint qadr = jnt_qposadr[jntadr] @@ -111,7 +111,7 @@ def _kinematics_level( # apply fixed translation and rotation relative to parent pid = body_parentid[bodyid] xpos = (xmat_in[worldid, pid] * body_pos[worldid, bodyid]) + xpos_in[worldid, pid] - xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[bodyid]) + xquat = math.mul_quat(xquat_in[worldid, pid], body_quat[worldid, bodyid]) for _ in range(jntnum): qadr = jnt_qposadr[jntadr] From a7111d5fa3e10cdf501e6f4384f57aaff1c296db Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 16:55:49 +0200 Subject: [PATCH 07/74] ipos --- 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 f088a1a7..e4ff460e 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -66,7 +66,7 @@ def _kinematics_level( body_jntadr: wp.array(dtype=int), body_pos: wp.array2d(dtype=wp.vec3), body_quat: wp.array2d(dtype=wp.quat), - body_ipos: wp.array(dtype=wp.vec3), + body_ipos: wp.array2d(dtype=wp.vec3), body_iquat: wp.array(dtype=wp.quat), jnt_type: wp.array(dtype=int), jnt_qposadr: wp.array(dtype=int), @@ -146,7 +146,7 @@ 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[bodyid], xquat) + xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid, bodyid], xquat) ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[bodyid])) From 018083ba26addeefb2150cec31e5d86308595700 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 16:57:56 +0200 Subject: [PATCH 08/74] body_mass --- mujoco_warp/_src/passive.py | 4 ++-- mujoco_warp/_src/sensor.py | 8 ++++---- mujoco_warp/_src/smooth.py | 20 ++++++++++---------- 3 files changed, 16 insertions(+), 16 deletions(-) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 1f82ec13..a6210bc8 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -119,7 +119,7 @@ def _gravity_force( opt_gravity: wp.vec3, body_parentid: wp.array(dtype=int), body_rootid: wp.array(dtype=int), - body_mass: wp.array(dtype=float), + body_mass: wp.array2d(dtype=float), body_gravcomp: wp.array(dtype=float), dof_bodyid: wp.array(dtype=int), # Data in: @@ -134,7 +134,7 @@ def _gravity_force( gravcomp = body_gravcomp[bodyid] if gravcomp: - force = -opt_gravity * body_mass[bodyid] * gravcomp + force = -opt_gravity * body_mass[worldid, bodyid] * gravcomp pos = xipos_in[worldid, bodyid] jac, _ = support.jac(body_parentid, body_rootid, dof_bodyid, subtree_com_in, cdof_in, pos, bodyid, dofid, worldid) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 5dde4a2c..e57a9fd5 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -279,7 +279,7 @@ def _frame_axis( @wp.func def _frame_quat( # Model: - body_iquat: wp.array(dtype=wp.quat), + body_iquat: wp.array2d(dtype=wp.quat), geom_bodyid: wp.array(dtype=int), geom_quat: wp.array(dtype=wp.quat), site_bodyid: wp.array(dtype=int), @@ -293,10 +293,10 @@ def _frame_quat( refid: int, ) -> wp.quat: if objtype == int(ObjType.BODY.value): - quat = math.mul_quat(xquat_in[worldid, objid], body_iquat[objid]) + quat = math.mul_quat(xquat_in[worldid, objid], body_iquat[worldid, objid]) if refid == -1: return quat - refquat = math.mul_quat(xquat_in[worldid, refid], body_iquat[refid]) + refquat = math.mul_quat(xquat_in[worldid, refid], body_iquat[worldid, refid]) elif objtype == int(ObjType.XBODY.value): quat = xquat_in[worldid, objid] if refid == -1: @@ -334,7 +334,7 @@ def _clock(time_in: wp.array(dtype=float), worldid: int) -> float: @wp.kernel def _sensor_pos( # Model: - body_iquat: wp.array(dtype=wp.quat), + body_iquat: wp.array2d(dtype=wp.quat), jnt_qposadr: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int), geom_quat: wp.array(dtype=wp.quat), diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index e4ff460e..29df672a 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -67,7 +67,7 @@ def _kinematics_level( body_pos: wp.array2d(dtype=wp.vec3), body_quat: wp.array2d(dtype=wp.quat), body_ipos: wp.array2d(dtype=wp.vec3), - body_iquat: wp.array(dtype=wp.quat), + body_iquat: wp.array2d(dtype=wp.quat), jnt_type: wp.array(dtype=int), jnt_qposadr: wp.array(dtype=int), jnt_pos: wp.array(dtype=wp.vec3), @@ -147,7 +147,7 @@ def _kinematics_level( 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, bodyid], xquat) - ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[bodyid])) + ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[worldid,bodyid])) @wp.kernel @@ -250,14 +250,14 @@ def kinematics(m: Model, d: Data): @wp.kernel def _subtree_com_init( # Model: - body_mass: wp.array(dtype=float), + body_mass: wp.array2d(dtype=float), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), # Data out: xipos_out: wp.array2d(dtype=wp.vec3), ): worldid, bodyid = wp.tid() - xipos_out[worldid, bodyid] = xipos_in[worldid, bodyid] * body_mass[bodyid] + xipos_out[worldid, bodyid] = xipos_in[worldid, bodyid] * body_mass[worldid, bodyid] @wp.kernel @@ -292,7 +292,7 @@ def _subtree_div( def _cinert( # Model: body_rootid: wp.array(dtype=int), - body_mass: wp.array(dtype=float), + body_mass: wp.array2d(dtype=float), body_inertia: wp.array(dtype=wp.vec3), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), @@ -304,7 +304,7 @@ def _cinert( worldid, bodyid = wp.tid() mat = ximat_in[worldid, bodyid] inert = body_inertia[bodyid] - mass = body_mass[bodyid] + mass = body_mass[worldid, bodyid] dif = xipos_in[worldid, bodyid] - subtree_com_in[worldid, body_rootid[bodyid]] # express inertia in com-based frame (mju_inertCom) @@ -1712,7 +1712,7 @@ def factor_solve_i(m, d, M, L, D, x, y): def _subtree_vel_forward( # Model: body_rootid: wp.array(dtype=int), - body_mass: wp.array(dtype=float), + body_mass: wp.array2d(dtype=float), body_inertia: wp.array(dtype=wp.vec3), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), @@ -1736,7 +1736,7 @@ def _subtree_vel_forward( # update linear velocity lin -= wp.cross(xipos - subtree_com_root, ang) - subtree_linvel_out[worldid, bodyid] = body_mass[bodyid] * lin + subtree_linvel_out[worldid, bodyid] = body_mass[worldid, bodyid] * lin dv = wp.transpose(ximat) @ ang dv[0] *= body_inertia[bodyid][0] dv[1] *= body_inertia[bodyid][1] @@ -1769,7 +1769,7 @@ def _linear_momentum( def _angular_momentum( # Model: body_parentid: wp.array(dtype=int), - body_mass: wp.array(dtype=float), + body_mass: wp.array2d(dtype=float), body_subtreemass: wp.array(dtype=float), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), @@ -1795,7 +1795,7 @@ def _angular_momentum( vel = subtree_bodyvel_in[worldid, bodyid] linvel = subtree_linvel_in[worldid, bodyid] linvel_parent = subtree_linvel_in[worldid, pid] - mass = body_mass[bodyid] + mass = body_mass[worldid, bodyid] subtreemass = body_subtreemass[bodyid] # momentum wrt body i From 4d18f8f9f855ebd4098fa2bad3f53216eb220d4a Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 16:58:32 +0200 Subject: [PATCH 09/74] subtreemass --- mujoco_warp/_src/smooth.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 29df672a..a7c3f06c 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -1749,7 +1749,7 @@ def _subtree_vel_forward( def _linear_momentum( # Model: body_parentid: wp.array(dtype=int), - body_subtreemass: wp.array(dtype=float), + body_subtreemass: wp.array2d(dtype=float), # Data in: subtree_linvel_in: wp.array2d(dtype=wp.vec3), # In: @@ -1762,7 +1762,7 @@ def _linear_momentum( if bodyid: pid = body_parentid[bodyid] wp.atomic_add(subtree_linvel_out[worldid], pid, subtree_linvel_in[worldid, bodyid]) - subtree_linvel_out[worldid, bodyid] /= wp.max(MJ_MINVAL, body_subtreemass[bodyid]) + subtree_linvel_out[worldid, bodyid] /= wp.max(MJ_MINVAL, body_subtreemass[worldid,bodyid]) @wp.kernel @@ -1770,7 +1770,7 @@ def _angular_momentum( # Model: body_parentid: wp.array(dtype=int), body_mass: wp.array2d(dtype=float), - body_subtreemass: wp.array(dtype=float), + body_subtreemass: wp.array2d(dtype=float), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3), @@ -1796,7 +1796,7 @@ def _angular_momentum( linvel = subtree_linvel_in[worldid, bodyid] linvel_parent = subtree_linvel_in[worldid, pid] mass = body_mass[worldid, bodyid] - subtreemass = body_subtreemass[bodyid] + subtreemass = body_subtreemass[worldid, bodyid] # momentum wrt body i dx = xipos - com From f2056ec11231ade0d8d5d6d9ded5f6a5e137fe7a Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 16:59:48 +0200 Subject: [PATCH 10/74] inertia --- mujoco_warp/_src/smooth.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index a7c3f06c..29071f23 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -293,7 +293,7 @@ def _cinert( # Model: body_rootid: wp.array(dtype=int), body_mass: wp.array2d(dtype=float), - body_inertia: wp.array(dtype=wp.vec3), + body_inertia: wp.array2d(dtype=wp.vec3), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), ximat_in: wp.array2d(dtype=wp.mat33), @@ -303,7 +303,7 @@ def _cinert( ): worldid, bodyid = wp.tid() mat = ximat_in[worldid, bodyid] - inert = body_inertia[bodyid] + inert = body_inertia[worldid,bodyid] mass = body_mass[worldid, bodyid] dif = xipos_in[worldid, bodyid] - subtree_com_in[worldid, body_rootid[bodyid]] # express inertia in com-based frame (mju_inertCom) @@ -1713,7 +1713,7 @@ def _subtree_vel_forward( # Model: body_rootid: wp.array(dtype=int), body_mass: wp.array2d(dtype=float), - body_inertia: wp.array(dtype=wp.vec3), + body_inertia: wp.array2d(dtype=wp.vec3), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), ximat_in: wp.array2d(dtype=wp.mat33), @@ -1738,9 +1738,9 @@ def _subtree_vel_forward( subtree_linvel_out[worldid, bodyid] = body_mass[worldid, bodyid] * lin dv = wp.transpose(ximat) @ ang - dv[0] *= body_inertia[bodyid][0] - dv[1] *= body_inertia[bodyid][1] - dv[2] *= body_inertia[bodyid][2] + dv[0] *= body_inertia[worldid, bodyid][0] + dv[1] *= body_inertia[worldid, bodyid][1] + dv[2] *= body_inertia[worldid, bodyid][2] subtree_angmom_out[worldid, bodyid] = ximat @ dv subtree_bodyvel_out[worldid, bodyid] = wp.spatial_vector(ang, lin) From a824ec32891c58f5d11a5e2247c072e0f4a349e3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:01:03 +0200 Subject: [PATCH 11/74] invweight0 --- mujoco_warp/_src/constraint.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index f193dfe7..9483bbd5 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -99,7 +99,7 @@ def _efc_equality_connect( opt_timestep: float, body_parentid: wp.array(dtype=int), body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=float), + body_invweight0: wp.array3d(dtype=float), dof_bodyid: wp.array(dtype=int), site_bodyid: wp.array(dtype=int), eq_obj1id: wp.array(dtype=int), @@ -195,7 +195,7 @@ def _efc_equality_connect( efc_J_out[efcid + 2, dofid] = j1mj2[2] Jqvel += j1mj2 * qvel_in[worldid, dofid] - invweight = body_invweight0[body1id, 0] + body_invweight0[body2id, 0] + invweight = body_invweight0[worldid, body1id, 0] + body_invweight0[worldid, body2id, 0] pos_imp = wp.length(pos) solref = eq_solref[i_eq] @@ -489,7 +489,7 @@ def _efc_equality_weld( opt_timestep: float, body_parentid: wp.array(dtype=int), body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=float), + body_invweight0: wp.array3d(dtype=float), dof_bodyid: wp.array(dtype=int), site_bodyid: wp.array(dtype=int), site_quat: wp.array(dtype=wp.quat), @@ -612,7 +612,7 @@ def _efc_equality_weld( crotq = math.mul_quat(quat1, quat) # copy axis components crot = wp.vec3(crotq[1], crotq[2], crotq[3]) * torquescale - invweight_t = body_invweight0[body1id, 0] + body_invweight0[body2id, 0] + invweight_t = body_invweight0[worldid, body1id, 0] + body_invweight0[worldid, body2id, 0] pos_imp = wp.sqrt(wp.length_sq(cpos) + wp.length_sq(crot)) @@ -641,7 +641,7 @@ def _efc_equality_weld( efc_frictionloss_out, ) - invweight_r = body_invweight0[body1id, 1] + body_invweight0[body2id, 1] + invweight_r = body_invweight0[worldid, body1id, 1] + body_invweight0[worldid, body2id, 1] for i in range(3): _update_efc_row( @@ -914,7 +914,7 @@ def _efc_contact_pyramidal( opt_impratio: float, body_parentid: wp.array(dtype=int), body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=float), + body_invweight0: wp.array3d(dtype=float), dof_bodyid: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int), # Data in: @@ -974,7 +974,7 @@ def _efc_contact_pyramidal( frame = frame_in[conid] # pyramidal has common invweight across all edges - invweight = body_invweight0[body1, 0] + body_invweight0[body2, 0] + invweight = body_invweight0[worldid, body1, 0] + body_invweight0[worldid, body2, 0] if condim > 1: dimid2 = dimid / 2 + 1 @@ -1060,7 +1060,7 @@ def _efc_contact_elliptic( opt_impratio: float, body_parentid: wp.array(dtype=int), body_rootid: wp.array(dtype=int), - body_invweight0: wp.array2d(dtype=float), + body_invweight0: wp.array3d(dtype=float), dof_bodyid: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int), # Data in: @@ -1158,7 +1158,7 @@ def _efc_contact_elliptic( efc_J_out[efcid, i] = J Jqvel += J * qvel_in[worldid, i] - invweight = body_invweight0[body1, 0] + body_invweight0[body2, 0] + invweight = body_invweight0[worldid, body1, 0] + body_invweight0[worldid, body2, 0] ref = solref_in[conid] pos_aref = pos From d8fa54223c26f01ff91278bd464ec55164cfa796 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:02:08 +0200 Subject: [PATCH 12/74] jnt_solref --- mujoco_warp/_src/constraint.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 9483bbd5..f42b75b1 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -672,7 +672,7 @@ def _efc_limit_slide_hinge( opt_timestep: float, jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), - jnt_solref: wp.array(dtype=wp.vec2), + jnt_solref: wp.array2d(dtype=wp.vec2), jnt_solimp: wp.array(dtype=vec5), jnt_range: wp.array2d(dtype=float), jnt_margin: wp.array(dtype=float), @@ -723,7 +723,7 @@ def _efc_limit_slide_hinge( pos, pos, dof_invweight0[dofadr], - jnt_solref[jntid], + jnt_solref[worldid, jntid], jnt_solimp[jntid], jntmargin, Jqvel, @@ -744,7 +744,7 @@ def _efc_limit_ball( opt_timestep: float, jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), - jnt_solref: wp.array(dtype=wp.vec2), + jnt_solref: wp.array2d(dtype=wp.vec2), jnt_solimp: wp.array(dtype=vec5), jnt_range: wp.array2d(dtype=float), jnt_margin: wp.array(dtype=float), @@ -803,7 +803,7 @@ def _efc_limit_ball( pos, pos, dof_invweight0[dofadr], - jnt_solref[jntid], + jnt_solref[worldid, jntid], jnt_solimp[jntid], jntmargin, Jqvel, From 0f07b2e784abedad852e7b54319e9b37ff0ee01e Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:02:47 +0200 Subject: [PATCH 13/74] jnt_solimp --- mujoco_warp/_src/constraint.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index f42b75b1..e5d3cbe4 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -673,7 +673,7 @@ def _efc_limit_slide_hinge( jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), jnt_solref: wp.array2d(dtype=wp.vec2), - jnt_solimp: wp.array(dtype=vec5), + jnt_solimp: wp.array2d(dtype=vec5), jnt_range: wp.array2d(dtype=float), jnt_margin: wp.array(dtype=float), jnt_limited_slide_hinge_adr: wp.array(dtype=int), @@ -724,7 +724,7 @@ def _efc_limit_slide_hinge( pos, dof_invweight0[dofadr], jnt_solref[worldid, jntid], - jnt_solimp[jntid], + jnt_solimp[worldid,jntid], jntmargin, Jqvel, 0.0, @@ -745,7 +745,7 @@ def _efc_limit_ball( jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), jnt_solref: wp.array2d(dtype=wp.vec2), - jnt_solimp: wp.array(dtype=vec5), + jnt_solimp: wp.array2d(dtype=vec5), jnt_range: wp.array2d(dtype=float), jnt_margin: wp.array(dtype=float), jnt_limited_ball_adr: wp.array(dtype=int), @@ -804,7 +804,7 @@ def _efc_limit_ball( pos, dof_invweight0[dofadr], jnt_solref[worldid, jntid], - jnt_solimp[jntid], + jnt_solimp[worldid, jntid], jntmargin, Jqvel, 0.0, From 738ee42bf5e514255b8793c8f84004e393d6947a Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:03:21 +0200 Subject: [PATCH 14/74] jnt_pos --- mujoco_warp/_src/smooth.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 29071f23..3ef7bcdc 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -70,7 +70,7 @@ def _kinematics_level( body_iquat: wp.array2d(dtype=wp.quat), jnt_type: wp.array(dtype=int), jnt_qposadr: wp.array(dtype=int), - jnt_pos: wp.array(dtype=wp.vec3), + jnt_pos: wp.array2d(dtype=wp.vec3), jnt_axis: wp.array(dtype=wp.vec3), # Data in: qpos_in: wp.array2d(dtype=float), @@ -117,7 +117,7 @@ def _kinematics_level( qadr = jnt_qposadr[jntadr] jnt_type_ = jnt_type[jntadr] jnt_axis_ = jnt_axis[jntadr] - xanchor = math.rot_vec_quat(jnt_pos[jntadr], xquat) + xpos + xanchor = math.rot_vec_quat(jnt_pos[worldid, jntadr], xquat) + xpos xaxis = math.rot_vec_quat(jnt_axis_, xquat) if jnt_type_ == wp.static(JointType.BALL.value): @@ -129,7 +129,7 @@ def _kinematics_level( ) xquat = math.mul_quat(xquat, qloc) # correct for off-center rotation - xpos = xanchor - math.rot_vec_quat(jnt_pos[jntadr], xquat) + xpos = xanchor - math.rot_vec_quat(jnt_pos[worldid, jntadr], xquat) elif jnt_type_ == wp.static(JointType.SLIDE.value): xpos += xaxis * (qpos[qadr] - qpos0[worldid, qadr]) elif jnt_type_ == wp.static(JointType.HINGE.value): @@ -137,7 +137,7 @@ def _kinematics_level( qloc_ = math.axis_angle_to_quat(jnt_axis_, qpos[qadr] - qpos0_) xquat = math.mul_quat(xquat, qloc_) # correct for off-center rotation - xpos = xanchor - math.rot_vec_quat(jnt_pos[jntadr], xquat) + xpos = xanchor - math.rot_vec_quat(jnt_pos[worldid, jntadr], xquat) xanchor_out[worldid, jntadr] = xanchor xaxis_out[worldid, jntadr] = xaxis From 2645791d6c1b3b4edf6962a3e75a2c5d61bee19e Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:03:47 +0200 Subject: [PATCH 15/74] jnt_stiffness --- mujoco_warp/_src/passive.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index a6210bc8..5ff68928 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -31,14 +31,14 @@ def _spring_passive( jnt_type: wp.array(dtype=int), jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), - jnt_stiffness: wp.array(dtype=float), + jnt_stiffness: wp.array2d(dtype=float), # Data in: qpos_in: wp.array2d(dtype=float), # Data out: qfrc_spring_out: wp.array2d(dtype=float), ): worldid, jntid = wp.tid() - stiffness = jnt_stiffness[jntid] + stiffness = jnt_stiffness[worldid,jntid] dofid = jnt_dofadr[jntid] if stiffness == 0.0: From 0e277b2bd3dc2090bcf20ecabbfe4e84507c9617 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:04:19 +0200 Subject: [PATCH 16/74] jnt_range --- mujoco_warp/_src/constraint.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index e5d3cbe4..984a5d40 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -674,7 +674,7 @@ def _efc_limit_slide_hinge( jnt_dofadr: wp.array(dtype=int), jnt_solref: wp.array2d(dtype=wp.vec2), jnt_solimp: wp.array2d(dtype=vec5), - jnt_range: wp.array2d(dtype=float), + jnt_range: wp.array3d(dtype=float), jnt_margin: wp.array(dtype=float), jnt_limited_slide_hinge_adr: wp.array(dtype=int), dof_invweight0: wp.array(dtype=float), @@ -697,7 +697,7 @@ def _efc_limit_slide_hinge( ): worldid, jntlimitedid = wp.tid() jntid = jnt_limited_slide_hinge_adr[jntlimitedid] - jntrange = jnt_range[jntid] + jntrange = jnt_range[worldid,jntid] qpos = qpos_in[worldid, jnt_qposadr[jntid]] jntmargin = jnt_margin[jntid] @@ -746,7 +746,7 @@ def _efc_limit_ball( jnt_dofadr: wp.array(dtype=int), jnt_solref: wp.array2d(dtype=wp.vec2), jnt_solimp: wp.array2d(dtype=vec5), - jnt_range: wp.array2d(dtype=float), + jnt_range: wp.array3d(dtype=float), jnt_margin: wp.array(dtype=float), jnt_limited_ball_adr: wp.array(dtype=int), dof_invweight0: wp.array(dtype=float), @@ -774,7 +774,7 @@ def _efc_limit_ball( qpos = qpos_in[worldid] jnt_quat = wp.quat(qpos[qposadr + 0], qpos[qposadr + 1], qpos[qposadr + 2], qpos[qposadr + 3]) axis_angle = math.quat_to_vel(jnt_quat) - jntrange = jnt_range[jntid] + jntrange = jnt_range[worldid, jntid] axis, angle = math.normalize_with_norm(axis_angle) jntmargin = jnt_margin[jntid] From f8a88f3e6ee6c13d44a1e87f4ad4dbba7fe13256 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:05:08 +0200 Subject: [PATCH 17/74] actfrcrange --- mujoco_warp/_src/forward.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index 0b104c75..f3980c9f 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -904,7 +904,7 @@ def _qfrc_actuator_sparse( nu: int, ngravcomp: int, jnt_actfrclimited: wp.array(dtype=bool), - jnt_actfrcrange: wp.array(dtype=wp.vec2), + jnt_actfrcrange: wp.array2d(dtype=wp.vec2), jnt_actgravcomp: wp.array(dtype=int), dof_jntid: wp.array(dtype=int), # Data in: @@ -928,7 +928,7 @@ def _qfrc_actuator_sparse( qfrc += qfrc_gravcomp_in[worldid, dofid] if jnt_actfrclimited[jntid]: - frcrange = jnt_actfrcrange[jntid] + frcrange = jnt_actfrcrange[worldid, jntid] qfrc = wp.clamp(qfrc, frcrange[0], frcrange[1]) qfrc_actuator_out[worldid, dofid] = qfrc @@ -939,7 +939,7 @@ def _qfrc_actuator_limited( # Model: ngravcomp: int, jnt_actfrclimited: wp.array(dtype=bool), - jnt_actfrcrange: wp.array(dtype=wp.vec2), + jnt_actfrcrange: wp.array2d(dtype=wp.vec2), jnt_actgravcomp: wp.array(dtype=int), dof_jntid: wp.array(dtype=int), # Data in: @@ -957,7 +957,7 @@ def _qfrc_actuator_limited( qfrc_dof += qfrc_gravcomp_in[worldid, dofid] if jnt_actfrclimited[jntid]: - frcrange = jnt_actfrcrange[jntid] + frcrange = jnt_actfrcrange[worldid, jntid] qfrc_dof = wp.clamp(qfrc_dof, frcrange[0], frcrange[1]) qfrc_actuator_out[worldid, dofid] = qfrc_dof From ba1dca48f5e9831d06922dc3e4a3682258777c28 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:05:54 +0200 Subject: [PATCH 18/74] jnt_margin --- mujoco_warp/_src/constraint.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 984a5d40..913aacf1 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -675,7 +675,7 @@ def _efc_limit_slide_hinge( jnt_solref: wp.array2d(dtype=wp.vec2), jnt_solimp: wp.array2d(dtype=vec5), jnt_range: wp.array3d(dtype=float), - jnt_margin: wp.array(dtype=float), + jnt_margin: wp.array2d(dtype=float), jnt_limited_slide_hinge_adr: wp.array(dtype=int), dof_invweight0: wp.array(dtype=float), # Data in: @@ -700,7 +700,7 @@ def _efc_limit_slide_hinge( jntrange = jnt_range[worldid,jntid] qpos = qpos_in[worldid, jnt_qposadr[jntid]] - jntmargin = jnt_margin[jntid] + jntmargin = jnt_margin[worldid, jntid] dist_min, dist_max = qpos - jntrange[0], jntrange[1] - qpos pos = wp.min(dist_min, dist_max) - jntmargin active = pos < 0 @@ -747,7 +747,7 @@ def _efc_limit_ball( jnt_solref: wp.array2d(dtype=wp.vec2), jnt_solimp: wp.array2d(dtype=vec5), jnt_range: wp.array3d(dtype=float), - jnt_margin: wp.array(dtype=float), + jnt_margin: wp.array2d(dtype=float), jnt_limited_ball_adr: wp.array(dtype=int), dof_invweight0: wp.array(dtype=float), # Data in: @@ -776,7 +776,7 @@ def _efc_limit_ball( axis_angle = math.quat_to_vel(jnt_quat) jntrange = jnt_range[worldid, jntid] axis, angle = math.normalize_with_norm(axis_angle) - jntmargin = jnt_margin[jntid] + jntmargin = jnt_margin[worldid, jntid] pos = wp.max(jntrange[0], jntrange[1]) - angle - jntmargin active = pos < 0 From 99ebf0b42664af1eec2431fdefd6641d2c9f3f93 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:07:28 +0200 Subject: [PATCH 19/74] dof_armature --- mujoco_warp/_src/smooth.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 3ef7bcdc..61642663 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -635,7 +635,7 @@ def _qM_sparse( dof_bodyid: wp.array(dtype=int), dof_parentid: wp.array(dtype=int), dof_Madr: wp.array(dtype=int), - dof_armature: wp.array(dtype=float), + dof_armature: wp.array2d(dtype=float), # Data in: cdof_in: wp.array2d(dtype=wp.spatial_vector), crb_in: wp.array2d(dtype=vec10), @@ -647,7 +647,7 @@ def _qM_sparse( bodyid = dof_bodyid[dofid] # init M(i,i) with armature inertia - qM_out[worldid, 0, madr_ij] = dof_armature[dofid] + qM_out[worldid, 0, madr_ij] = dof_armature[worldid,dofid] # precompute buf = crb_body_i * cdof_i buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid]) @@ -664,7 +664,7 @@ def _qM_dense( # Model: dof_bodyid: wp.array(dtype=int), dof_parentid: wp.array(dtype=int), - dof_armature: wp.array(dtype=float), + dof_armature: wp.array2d(dtype=float), # Data in: cdof_in: wp.array2d(dtype=wp.spatial_vector), crb_in: wp.array2d(dtype=vec10), @@ -674,7 +674,7 @@ def _qM_dense( worldid, dofid = wp.tid() bodyid = dof_bodyid[dofid] # init M(i,i) with armature inertia - M = dof_armature[dofid] + M = dof_armature[worldid, dofid] # precompute buf = crb_body_i * cdof_i buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid]) From 4e2a4055381c2e6083f73370037f1954870a1d74 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:08:48 +0200 Subject: [PATCH 20/74] dof_damping --- mujoco_warp/_src/forward.py | 12 ++++++------ mujoco_warp/_src/passive.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index f3980c9f..a851a406 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -267,7 +267,7 @@ def _euler_damp_qfrc_sparse( # Model: opt_timestep: float, dof_Madr: wp.array(dtype=int), - dof_damping: wp.array(dtype=float), + dof_damping: wp.array2d(dtype=float), # Data in: qfrc_smooth_in: wp.array2d(dtype=float), qfrc_constraint_in: wp.array2d(dtype=float), @@ -278,7 +278,7 @@ def _euler_damp_qfrc_sparse( worldid, tid = wp.tid() adr = dof_Madr[tid] - qM_integration_out[worldid, 0, adr] += opt_timestep * dof_damping[tid] + qM_integration_out[worldid, 0, adr] += opt_timestep * dof_damping[worldid,tid] qfrc_integration_out[worldid, tid] = qfrc_smooth_in[worldid, tid] + qfrc_constraint_in[worldid, tid] @@ -314,7 +314,7 @@ def _tile_euler_dense(tile: TileSet): @nested_kernel def euler_dense( # Model: - dof_damping: wp.array(dtype=float), + dof_damping: wp.array2d(dtype=float), opt_timestep: float, # Data in: qM_in: wp.array3d(dtype=float), @@ -330,7 +330,7 @@ def euler_dense( dofid = adr_in[nodeid] M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid)) - damping_tile = wp.tile_load(dof_damping, shape=(TILE_SIZE,), offset=(dofid,)) + damping_tile = wp.tile_load(dof_damping[worldid], shape=(TILE_SIZE,), offset=(dofid,)) damping_scaled = damping_tile * opt_timestep qm_integration_tile = wp.tile_diag_add(M_tile, damping_scaled) @@ -523,7 +523,7 @@ def subtract_multiply(x: float, y: float): @nested_kernel def implicit_actuator_qderiv( # Model: - dof_damping: wp.array(dtype=float), + dof_damping: wp.array2d(dtype=float), # Data in: actuator_moment_in: wp.array3d(dtype=float), qM_in: wp.array3d(dtype=float), @@ -562,7 +562,7 @@ def implicit_actuator_qderiv( qderiv_tile = wp.tile_zeros(shape=(TILE_NV_SIZE, TILE_NV_SIZE), dtype=wp.float32) if wp.static(passive_enabled): - dof_damping_tile = wp.tile_load(dof_damping, shape=TILE_NV_SIZE, offset=offset_nv) + dof_damping_tile = wp.tile_load(dof_damping[worldid], shape=TILE_NV_SIZE, offset=offset_nv) negative = wp.neg(dof_damping_tile) qderiv_tile = wp.tile_diag_add(qderiv_tile, negative) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 5ff68928..52f84204 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -97,7 +97,7 @@ def _spring_passive( @wp.kernel def _damper_passive( # Model: - dof_damping: wp.array(dtype=float), + dof_damping: wp.array2d(dtype=float), # Data in: qvel_in: wp.array2d(dtype=float), qfrc_spring_in: wp.array2d(dtype=float), @@ -107,7 +107,7 @@ def _damper_passive( ): worldid, dofid = wp.tid() - qfrc_damper = -dof_damping[dofid] * qvel_in[worldid, dofid] + qfrc_damper = -dof_damping[worldid, dofid] * qvel_in[worldid, dofid] qfrc_damper_out[worldid, dofid] = qfrc_damper qfrc_passive_out[worldid, dofid] = qfrc_damper + qfrc_spring_in[worldid, dofid] From d9800a599ed4b2d65a3a4ea4fb042330de089876 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:09:48 +0200 Subject: [PATCH 21/74] dof_invweight' --- mujoco_warp/_src/constraint.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 913aacf1..03a3320f 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -234,7 +234,7 @@ def _efc_equality_joint( qpos0: wp.array2d(dtype=float), jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), - dof_invweight0: wp.array(dtype=float), + dof_invweight0: wp.array2d(dtype=float), eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), eq_solref: wp.array(dtype=wp.vec2), @@ -289,14 +289,14 @@ def _efc_equality_joint( pos = qpos_in[worldid, qposadr1] - qpos0[worldid,qposadr1] - rhs Jqvel = qvel_in[worldid, dofadr1] - qvel_in[worldid, dofadr2] * deriv_2 - invweight = dof_invweight0[dofadr1] + dof_invweight0[dofadr2] + invweight = dof_invweight0[worldid, dofadr1] + dof_invweight0[worldid, dofadr2] efc_J_out[efcid, dofadr2] = -deriv_2 else: # Single joint constraint pos = qpos_in[worldid, qposadr1] - qpos0[worldid, qposadr1] - data[0] Jqvel = qvel_in[worldid, dofadr1] - invweight = dof_invweight0[dofadr1] + invweight = dof_invweight0[worldid, dofadr1] # Update constraint parameters _update_efc_row( @@ -426,7 +426,7 @@ def _efc_equality_tendon( def _efc_friction( # Model: opt_timestep: float, - dof_invweight0: wp.array(dtype=float), + dof_invweight0: wp.array2d(dtype=float), dof_frictionloss: wp.array(dtype=float), dof_solimp: wp.array(dtype=vec5), dof_solref: wp.array(dtype=wp.vec2), @@ -465,7 +465,7 @@ def _efc_friction( efcid, 0.0, 0.0, - dof_invweight0[dofid], + dof_invweight0[worldid, dofid], dof_solref[dofid], dof_solimp[dofid], 0.0, @@ -677,7 +677,7 @@ def _efc_limit_slide_hinge( jnt_range: wp.array3d(dtype=float), jnt_margin: wp.array2d(dtype=float), jnt_limited_slide_hinge_adr: wp.array(dtype=int), - dof_invweight0: wp.array(dtype=float), + dof_invweight0: wp.array2d(dtype=float), # Data in: nefc_in: wp.array(dtype=int), qpos_in: wp.array2d(dtype=float), @@ -722,7 +722,7 @@ def _efc_limit_slide_hinge( efcid, pos, pos, - dof_invweight0[dofadr], + dof_invweight0[worldid, dofadr], jnt_solref[worldid, jntid], jnt_solimp[worldid,jntid], jntmargin, @@ -749,7 +749,7 @@ def _efc_limit_ball( jnt_range: wp.array3d(dtype=float), jnt_margin: wp.array2d(dtype=float), jnt_limited_ball_adr: wp.array(dtype=int), - dof_invweight0: wp.array(dtype=float), + dof_invweight0: wp.array2d(dtype=float), # Data in: nefc_in: wp.array(dtype=int), qpos_in: wp.array2d(dtype=float), @@ -802,7 +802,7 @@ def _efc_limit_ball( efcid, pos, pos, - dof_invweight0[dofadr], + dof_invweight0[worldid, dofadr], jnt_solref[worldid, jntid], jnt_solimp[worldid, jntid], jntmargin, From ae0c11d7d65d5a10bf7eae44ce9853f53f08e307 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:10:26 +0200 Subject: [PATCH 22/74] dof_frictionloss --- mujoco_warp/_src/constraint.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 03a3320f..76204829 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -427,7 +427,7 @@ def _efc_friction( # Model: opt_timestep: float, dof_invweight0: wp.array2d(dtype=float), - dof_frictionloss: wp.array(dtype=float), + dof_frictionloss: wp.array2d(dtype=float), dof_solimp: wp.array(dtype=vec5), dof_solref: wp.array(dtype=wp.vec2), # Data in: @@ -449,7 +449,7 @@ def _efc_friction( # TODO(team): tendon worldid, dofid = wp.tid() - if dof_frictionloss[dofid] <= 0.0: + if dof_frictionloss[worldid, dofid] <= 0.0: return efcid = wp.atomic_add(nefc_out, 0, 1) @@ -470,7 +470,7 @@ def _efc_friction( dof_solimp[dofid], 0.0, Jqvel, - dof_frictionloss[dofid], + dof_frictionloss[worldid, dofid], dofid, efc_id_out, efc_pos_out, From 8ac6b916f8ebd719013f8caa2a7604d9027e95dd Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:10:50 +0200 Subject: [PATCH 23/74] dof_solimp --- mujoco_warp/_src/constraint.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 76204829..5e3b87e4 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -428,7 +428,7 @@ def _efc_friction( opt_timestep: float, dof_invweight0: wp.array2d(dtype=float), dof_frictionloss: wp.array2d(dtype=float), - dof_solimp: wp.array(dtype=vec5), + dof_solimp: wp.array2d(dtype=vec5), dof_solref: wp.array(dtype=wp.vec2), # Data in: qvel_in: wp.array2d(dtype=float), @@ -467,7 +467,7 @@ def _efc_friction( 0.0, dof_invweight0[worldid, dofid], dof_solref[dofid], - dof_solimp[dofid], + dof_solimp[worldid, dofid], 0.0, Jqvel, dof_frictionloss[worldid, dofid], From ad3436b9af8dbb56cdcadea719a57676ed8c7449 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:11:33 +0200 Subject: [PATCH 24/74] dof_solref --- mujoco_warp/_src/constraint.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 5e3b87e4..7a15f700 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -429,7 +429,7 @@ def _efc_friction( dof_invweight0: wp.array2d(dtype=float), dof_frictionloss: wp.array2d(dtype=float), dof_solimp: wp.array2d(dtype=vec5), - dof_solref: wp.array(dtype=wp.vec2), + dof_solref: wp.array2d(dtype=wp.vec2), # Data in: qvel_in: wp.array2d(dtype=float), # In: @@ -466,7 +466,7 @@ def _efc_friction( 0.0, 0.0, dof_invweight0[worldid, dofid], - dof_solref[dofid], + dof_solref[worldid, dofid], dof_solimp[worldid, dofid], 0.0, Jqvel, From c685f0342d589826f4ce6b61950ae03944430433 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:12:38 +0200 Subject: [PATCH 25/74] geom_priority --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 562cf6ff..4001b4ac 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -191,7 +191,7 @@ def _box_box( # Model: geom_type: wp.array(dtype=int), geom_condim: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), + geom_priority: wp.array2d(dtype=int), geom_solmix: wp.array(dtype=float), geom_solref: wp.array(dtype=wp.vec2), geom_solimp: wp.array(dtype=vec5), @@ -245,7 +245,7 @@ def _box_box( geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, - geom_priority, + geom_priority[worldid], geom_solmix, geom_solref, geom_solimp, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 4c735062..c093c5e6 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -708,7 +708,7 @@ def gjk_epa_sparse( geom_type: wp.array(dtype=int), geom_condim: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), + geom_priority: wp.array2d(dtype=int), geom_solmix: wp.array(dtype=float), geom_solref: wp.array(dtype=wp.vec2), geom_solimp: wp.array(dtype=vec5), @@ -755,7 +755,7 @@ def gjk_epa_sparse( worldid = collision_worldid_in[tid] geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, - geom_priority, + geom_priority[worldid], geom_solmix, geom_solref, geom_solimp, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 0bffc3a4..d69a2d57 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1724,7 +1724,7 @@ def _primitive_narrowphase( geom_type: wp.array(dtype=int), geom_condim: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), + geom_priority: wp.array2d(dtype=int), geom_solmix: wp.array(dtype=float), geom_solref: wp.array(dtype=wp.vec2), geom_solimp: wp.array(dtype=vec5), @@ -1770,7 +1770,7 @@ def _primitive_narrowphase( geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, - geom_priority, + geom_priority[worldid], geom_solmix, geom_solref, geom_solimp, From f814ff9244347e3130ce518bfcf6ff04c96ef01d Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:13:32 +0200 Subject: [PATCH 26/74] geom_solmix --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 4001b4ac..145e4299 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -192,7 +192,7 @@ def _box_box( geom_type: wp.array(dtype=int), geom_condim: wp.array(dtype=int), geom_priority: wp.array2d(dtype=int), - geom_solmix: wp.array(dtype=float), + geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array(dtype=wp.vec2), geom_solimp: wp.array(dtype=vec5), geom_size: wp.array(dtype=wp.vec3), @@ -246,7 +246,7 @@ def _box_box( geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, geom_priority[worldid], - geom_solmix, + geom_solmix[worldid], geom_solref, geom_solimp, geom_friction, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index c093c5e6..2684afef 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -709,7 +709,7 @@ def gjk_epa_sparse( geom_condim: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), geom_priority: wp.array2d(dtype=int), - geom_solmix: wp.array(dtype=float), + geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array(dtype=wp.vec2), geom_solimp: wp.array(dtype=vec5), geom_size: wp.array(dtype=wp.vec3), @@ -756,7 +756,7 @@ def gjk_epa_sparse( geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, geom_priority[worldid], - geom_solmix, + geom_solmix[worldid], geom_solref, geom_solimp, geom_friction, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index d69a2d57..a1fc32b6 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1725,7 +1725,7 @@ def _primitive_narrowphase( geom_condim: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), geom_priority: wp.array2d(dtype=int), - geom_solmix: wp.array(dtype=float), + geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array(dtype=wp.vec2), geom_solimp: wp.array(dtype=vec5), geom_size: wp.array(dtype=wp.vec3), @@ -1771,7 +1771,7 @@ def _primitive_narrowphase( geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, geom_priority[worldid], - geom_solmix, + geom_solmix[worldid], geom_solref, geom_solimp, geom_friction, From 1a3a8be3ff64b93efc7ffd8d9b68f03b8897d386 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:14:26 +0200 Subject: [PATCH 27/74] geom_solref --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 145e4299..9fda659d 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -193,7 +193,7 @@ def _box_box( geom_condim: wp.array(dtype=int), geom_priority: wp.array2d(dtype=int), geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array(dtype=wp.vec2), + geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array(dtype=vec5), geom_size: wp.array(dtype=wp.vec3), geom_friction: wp.array(dtype=wp.vec3), @@ -247,7 +247,7 @@ def _box_box( geom_condim, geom_priority[worldid], geom_solmix[worldid], - geom_solref, + geom_solref[worldid], geom_solimp, geom_friction, geom_margin, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 2684afef..f510726a 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -710,7 +710,7 @@ def gjk_epa_sparse( geom_dataid: wp.array(dtype=int), geom_priority: wp.array2d(dtype=int), geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array(dtype=wp.vec2), + geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array(dtype=vec5), geom_size: wp.array(dtype=wp.vec3), geom_friction: wp.array(dtype=wp.vec3), @@ -757,7 +757,7 @@ def gjk_epa_sparse( geom_condim, geom_priority[worldid], geom_solmix[worldid], - geom_solref, + geom_solref[worldid], geom_solimp, geom_friction, geom_margin, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index a1fc32b6..8e07b0bc 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1726,7 +1726,7 @@ def _primitive_narrowphase( geom_dataid: wp.array(dtype=int), geom_priority: wp.array2d(dtype=int), geom_solmix: wp.array2d(dtype=float), - geom_solref: wp.array(dtype=wp.vec2), + geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array(dtype=vec5), geom_size: wp.array(dtype=wp.vec3), geom_friction: wp.array(dtype=wp.vec3), @@ -1772,7 +1772,7 @@ def _primitive_narrowphase( geom_condim, geom_priority[worldid], geom_solmix[worldid], - geom_solref, + geom_solref[worldid], geom_solimp, geom_friction, geom_margin, From d8762ad3321da9e738097641ef7f42470ce47dbb Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:15:38 +0200 Subject: [PATCH 28/74] geom_solimp --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 9fda659d..49214452 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -194,7 +194,7 @@ def _box_box( geom_priority: wp.array2d(dtype=int), geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array(dtype=vec5), + geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array(dtype=wp.vec3), geom_friction: wp.array(dtype=wp.vec3), geom_margin: wp.array(dtype=float), @@ -248,7 +248,7 @@ def _box_box( geom_priority[worldid], geom_solmix[worldid], geom_solref[worldid], - geom_solimp, + geom_solimp[worldid], geom_friction, geom_margin, geom_gap, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index f510726a..82bcb162 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -711,7 +711,7 @@ def gjk_epa_sparse( geom_priority: wp.array2d(dtype=int), geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array(dtype=vec5), + geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array(dtype=wp.vec3), geom_friction: wp.array(dtype=wp.vec3), geom_margin: wp.array(dtype=float), @@ -758,7 +758,7 @@ def gjk_epa_sparse( geom_priority[worldid], geom_solmix[worldid], geom_solref[worldid], - geom_solimp, + geom_solimp[worldid], geom_friction, geom_margin, geom_gap, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 8e07b0bc..005c4a12 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1727,7 +1727,7 @@ def _primitive_narrowphase( geom_priority: wp.array2d(dtype=int), geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), - geom_solimp: wp.array(dtype=vec5), + geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array(dtype=wp.vec3), geom_friction: wp.array(dtype=wp.vec3), geom_margin: wp.array(dtype=float), @@ -1773,7 +1773,7 @@ def _primitive_narrowphase( geom_priority[worldid], geom_solmix[worldid], geom_solref[worldid], - geom_solimp, + geom_solimp[worldid], geom_friction, geom_margin, geom_gap, From 57e550cee716f402bb3a193202b7b636f9d706f6 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:17:04 +0200 Subject: [PATCH 29/74] geom_size --- mujoco_warp/_src/collision_box.py | 6 +++--- mujoco_warp/_src/collision_convex.py | 6 +++--- mujoco_warp/_src/collision_primitive.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 49214452..875cd0ab 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -195,7 +195,7 @@ def _box_box( geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array(dtype=wp.vec3), + geom_size: wp.array2d(dtype=wp.vec3), geom_friction: wp.array(dtype=wp.vec3), geom_margin: wp.array(dtype=float), geom_gap: wp.array(dtype=float), @@ -271,8 +271,8 @@ def _box_box( trans_atob = b_mat_inv @ (a_pos - b_pos) rot_atob = b_mat_inv @ a_mat - a_size = geom_size[ga] - b_size = geom_size[gb] + a_size = geom_size[worldid, ga] + b_size = geom_size[worldid, gb] a = box(rot_atob, trans_atob, a_size) b = box(wp.identity(3, wp.float32), wp.vec3(0.0), b_size) diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 82bcb162..4385c7b6 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -712,7 +712,7 @@ def gjk_epa_sparse( geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array(dtype=wp.vec3), + geom_size: wp.array2d(dtype=wp.vec3), geom_friction: wp.array(dtype=wp.vec3), geom_margin: wp.array(dtype=float), geom_gap: wp.array(dtype=float), @@ -782,7 +782,7 @@ def gjk_epa_sparse( geom1 = _geom( geom_dataid, - geom_size, + geom_size[worldid], mesh_vertadr, mesh_vertnum, geom_xpos_in, @@ -793,7 +793,7 @@ def gjk_epa_sparse( geom2 = _geom( geom_dataid, - geom_size, + geom_size[worldid], mesh_vertadr, mesh_vertnum, geom_xpos_in, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 005c4a12..0a0f26d4 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1728,7 +1728,7 @@ def _primitive_narrowphase( geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), - geom_size: wp.array(dtype=wp.vec3), + geom_size: wp.array2d(dtype=wp.vec3), geom_friction: wp.array(dtype=wp.vec3), geom_margin: wp.array(dtype=float), geom_gap: wp.array(dtype=float), @@ -1795,7 +1795,7 @@ def _primitive_narrowphase( geom1 = _geom( geom_dataid, - geom_size, + geom_size[worldid], mesh_vertadr, mesh_vertnum, geom_xpos_in, @@ -1805,7 +1805,7 @@ def _primitive_narrowphase( ) geom2 = _geom( geom_dataid, - geom_size, + geom_size[worldid], mesh_vertadr, mesh_vertnum, geom_xpos_in, From 83ac91af30e503efa2b1403016d757e60fa8994c Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:18:49 +0200 Subject: [PATCH 30/74] geom_rbound --- mujoco_warp/_src/collision_driver.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index 799051cc..bb9fcdd6 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -124,7 +124,7 @@ def _upper_tri_index(n: int, i: int, j: int) -> int: @wp.kernel def _sap_project( # Model: - geom_rbound: wp.array(dtype=float), + geom_rbound: wp.array2d(dtype=float), geom_margin: wp.array(dtype=float), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -138,7 +138,7 @@ def _sap_project( worldid, geomid = wp.tid() xpos = geom_xpos_in[worldid, geomid] - rbound = geom_rbound[geomid] + rbound = geom_rbound[worldid,geomid] if rbound == 0.0: # geom is a plane @@ -182,7 +182,7 @@ def _sap_broadphase( # Model: ngeom: int, geom_type: wp.array(dtype=int), - geom_rbound: wp.array(dtype=float), + geom_rbound: wp.array2d(dtype=float), geom_margin: wp.array(dtype=float), nxn_pairid: wp.array(dtype=int), # Data in: @@ -233,7 +233,7 @@ def _sap_broadphase( continue if _sphere_filter( - geom_rbound, + geom_rbound[worldid], geom_margin, geom_xpos_in, geom_xmat_in, @@ -343,7 +343,7 @@ def sap_broadphase(m: Model, d: Data): def _nxn_broadphase( # Model: geom_type: wp.array(dtype=int), - geom_rbound: wp.array(dtype=float), + geom_rbound: wp.array2d(dtype=float), geom_margin: wp.array(dtype=float), nxn_geom_pair: wp.array(dtype=wp.vec2i), nxn_pairid: wp.array(dtype=int), @@ -368,7 +368,7 @@ def _nxn_broadphase( geom2 = geom[1] if _sphere_filter( - geom_rbound, + geom_rbound[worldid], geom_margin, geom_xpos_in, geom_xmat_in, From 69dcf51bc5f0c1313e0822f83b105833020999d8 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:19:16 +0200 Subject: [PATCH 31/74] geom_pos --- 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 61642663..9cc2fbfd 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -154,7 +154,7 @@ def _kinematics_level( def _geom_local_to_global( # Model: geom_bodyid: wp.array(dtype=int), - geom_pos: wp.array(dtype=wp.vec3), + geom_pos: wp.array2d(dtype=wp.vec3), geom_quat: wp.array(dtype=wp.quat), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), @@ -167,7 +167,7 @@ def _geom_local_to_global( bodyid = geom_bodyid[geomid] xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] - geom_xpos_out[worldid, geomid] = xpos + math.rot_vec_quat(geom_pos[geomid], xquat) + geom_xpos_out[worldid, geomid] = xpos + math.rot_vec_quat(geom_pos[worldid,geomid], xquat) geom_xmat_out[worldid, geomid] = math.quat_to_mat(math.mul_quat(xquat, geom_quat[geomid])) From f1b85f9a7a14fe3ffcc2d17d9b61354cf95027cc Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:20:50 +0200 Subject: [PATCH 32/74] geom_quat --- mujoco_warp/_src/sensor.py | 6 +++--- mujoco_warp/_src/smooth.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index e57a9fd5..64afc290 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -281,7 +281,7 @@ def _frame_quat( # Model: body_iquat: wp.array2d(dtype=wp.quat), geom_bodyid: wp.array(dtype=int), - geom_quat: wp.array(dtype=wp.quat), + geom_quat: wp.array2d(dtype=wp.quat), site_bodyid: wp.array(dtype=int), site_quat: wp.array(dtype=wp.quat), # Data in: @@ -303,10 +303,10 @@ def _frame_quat( return quat refquat = xquat_in[worldid, refid] elif objtype == int(ObjType.GEOM.value): - quat = math.mul_quat(xquat_in[worldid, geom_bodyid[objid]], geom_quat[objid]) + quat = math.mul_quat(xquat_in[worldid, geom_bodyid[objid]], geom_quat[worldid, objid]) if refid == -1: return quat - refquat = math.mul_quat(xquat_in[worldid, geom_bodyid[refid]], geom_quat[refid]) + refquat = math.mul_quat(xquat_in[worldid, geom_bodyid[refid]], geom_quat[worldid,refid]) elif objtype == int(ObjType.SITE.value): quat = math.mul_quat(xquat_in[worldid, site_bodyid[objid]], site_quat[objid]) if refid == -1: diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 9cc2fbfd..13891c73 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -155,7 +155,7 @@ def _geom_local_to_global( # Model: geom_bodyid: wp.array(dtype=int), geom_pos: wp.array2d(dtype=wp.vec3), - geom_quat: wp.array(dtype=wp.quat), + geom_quat: wp.array2d(dtype=wp.quat), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), @@ -168,7 +168,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,geomid], xquat) - geom_xmat_out[worldid, geomid] = math.quat_to_mat(math.mul_quat(xquat, geom_quat[geomid])) + geom_xmat_out[worldid, geomid] = math.quat_to_mat(math.mul_quat(xquat, geom_quat[worldid, geomid])) @wp.kernel From 8f4fa7c5a2c94f33dfc75efb03dce3f4061adc1a Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:21:43 +0200 Subject: [PATCH 33/74] geom_friction --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 875cd0ab..39c10b86 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -196,7 +196,7 @@ def _box_box( geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array2d(dtype=wp.vec3), - geom_friction: wp.array(dtype=wp.vec3), + geom_friction: wp.array2d(dtype=wp.vec3), geom_margin: wp.array(dtype=float), geom_gap: wp.array(dtype=float), pair_dim: wp.array(dtype=int), @@ -249,7 +249,7 @@ def _box_box( geom_solmix[worldid], geom_solref[worldid], geom_solimp[worldid], - geom_friction, + geom_friction[worldid], geom_margin, geom_gap, pair_dim, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 4385c7b6..c37a24cc 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -713,7 +713,7 @@ def gjk_epa_sparse( geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array2d(dtype=wp.vec3), - geom_friction: wp.array(dtype=wp.vec3), + geom_friction: wp.array2d(dtype=wp.vec3), geom_margin: wp.array(dtype=float), geom_gap: wp.array(dtype=float), mesh_vertadr: wp.array(dtype=int), @@ -759,7 +759,7 @@ def gjk_epa_sparse( geom_solmix[worldid], geom_solref[worldid], geom_solimp[worldid], - geom_friction, + geom_friction[worldid], geom_margin, geom_gap, pair_dim, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 0a0f26d4..f80dcf4c 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1729,7 +1729,7 @@ def _primitive_narrowphase( geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array2d(dtype=wp.vec3), - geom_friction: wp.array(dtype=wp.vec3), + geom_friction: wp.array2d(dtype=wp.vec3), geom_margin: wp.array(dtype=float), geom_gap: wp.array(dtype=float), mesh_vertadr: wp.array(dtype=int), @@ -1774,7 +1774,7 @@ def _primitive_narrowphase( geom_solmix[worldid], geom_solref[worldid], geom_solimp[worldid], - geom_friction, + geom_friction[worldid], geom_margin, geom_gap, pair_dim, From 6de0c311352218b51f582bb29bb50f081b17ca0b Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:24:14 +0200 Subject: [PATCH 34/74] geom_margin --- mujoco_warp/_src/collision_box.py | 6 +++--- mujoco_warp/_src/collision_convex.py | 6 +++--- mujoco_warp/_src/collision_driver.py | 24 ++++++++++++------------ mujoco_warp/_src/collision_primitive.py | 2 +- 4 files changed, 19 insertions(+), 19 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 39c10b86..ada1d52f 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -197,7 +197,7 @@ def _box_box( geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array(dtype=float), + geom_margin: wp.array2d(dtype=float), geom_gap: wp.array(dtype=float), pair_dim: wp.array(dtype=int), pair_solref: wp.array(dtype=wp.vec2), @@ -250,7 +250,7 @@ def _box_box( geom_solref[worldid], geom_solimp[worldid], geom_friction[worldid], - geom_margin, + geom_margin[worldid], geom_gap, pair_dim, pair_solref, @@ -363,7 +363,7 @@ def _box_box( for i in range(4): pos[i] = pos[idx] - margin = wp.max(geom_margin[ga], geom_margin[gb]) + margin = wp.max(geom_margin[worldid, ga], geom_margin[worldid, gb]) for i in range(4): pos_glob = b_mat @ pos[i] + b_pos n_glob = b_mat @ sep_axis diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index c37a24cc..71467bad 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -714,7 +714,7 @@ def gjk_epa_sparse( geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array(dtype=float), + geom_margin: wp.array2d(dtype=float), geom_gap: wp.array(dtype=float), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), @@ -760,7 +760,7 @@ def gjk_epa_sparse( geom_solref[worldid], geom_solimp[worldid], geom_friction[worldid], - geom_margin, + geom_margin[worldid], geom_gap, pair_dim, pair_solref, @@ -802,7 +802,7 @@ def gjk_epa_sparse( g2, ) - margin = wp.max(geom_margin[g1], geom_margin[g2]) + margin = wp.max(geom_margin[worldid, g1], geom_margin[worldid, g2]) simplex, normal = _gjk(mesh_vert, geom1, geom2) diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index bb9fcdd6..26f16d96 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -32,8 +32,8 @@ @wp.func def _sphere_filter( # Model: - geom_rbound: wp.array(dtype=float), - geom_margin: wp.array(dtype=float), + geom_rbound: wp.array2d(dtype=float), + geom_margin: wp.array2d(dtype=float), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), geom_xmat_in: wp.array2d(dtype=wp.mat33), @@ -42,12 +42,12 @@ def _sphere_filter( geom2: int, worldid: int, ) -> bool: - margin1 = geom_margin[geom1] - margin2 = geom_margin[geom2] + margin1 = geom_margin[worldid, geom1] + margin2 = geom_margin[worldid, geom2] pos1 = geom_xpos_in[worldid, geom1] pos2 = geom_xpos_in[worldid, geom2] - size1 = geom_rbound[geom1] - size2 = geom_rbound[geom2] + size1 = geom_rbound[worldid, geom1] + size2 = geom_rbound[worldid, geom2] bound = size1 + size2 + wp.max(margin1, margin2) dif = pos2 - pos1 @@ -125,7 +125,7 @@ def _upper_tri_index(n: int, i: int, j: int) -> int: def _sap_project( # Model: geom_rbound: wp.array2d(dtype=float), - geom_margin: wp.array(dtype=float), + geom_margin: wp.array2d(dtype=float), # Data in: geom_xpos_in: wp.array2d(dtype=wp.vec3), # In: @@ -144,7 +144,7 @@ def _sap_project( # geom is a plane rbound = MJ_MAXVAL - radius = rbound + geom_margin[geomid] + radius = rbound + geom_margin[worldid,geomid] center = wp.dot(direction_in, xpos) sap_projection_lower_out[worldid, geomid] = center - radius @@ -183,7 +183,7 @@ def _sap_broadphase( ngeom: int, geom_type: wp.array(dtype=int), geom_rbound: wp.array2d(dtype=float), - geom_margin: wp.array(dtype=float), + geom_margin: wp.array2d(dtype=float), nxn_pairid: wp.array(dtype=int), # Data in: nworld_in: int, @@ -233,7 +233,7 @@ def _sap_broadphase( continue if _sphere_filter( - geom_rbound[worldid], + geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, @@ -344,7 +344,7 @@ def _nxn_broadphase( # Model: geom_type: wp.array(dtype=int), geom_rbound: wp.array2d(dtype=float), - geom_margin: wp.array(dtype=float), + geom_margin: wp.array2d(dtype=float), nxn_geom_pair: wp.array(dtype=wp.vec2i), nxn_pairid: wp.array(dtype=int), # Data in: @@ -368,7 +368,7 @@ def _nxn_broadphase( geom2 = geom[1] if _sphere_filter( - geom_rbound[worldid], + geom_rbound, geom_margin, geom_xpos_in, geom_xmat_in, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index f80dcf4c..b7444661 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1775,7 +1775,7 @@ def _primitive_narrowphase( geom_solref[worldid], geom_solimp[worldid], geom_friction[worldid], - geom_margin, + geom_margin[worldid], geom_gap, pair_dim, pair_solref, From fe4ae67c106909c026613ab8d1ab58e4323f9881 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:25:10 +0200 Subject: [PATCH 35/74] geom_gap --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index ada1d52f..85b27fda 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -198,7 +198,7 @@ def _box_box( geom_size: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d(dtype=wp.vec3), geom_margin: wp.array2d(dtype=float), - geom_gap: wp.array(dtype=float), + geom_gap: wp.array2d(dtype=float), pair_dim: wp.array(dtype=int), pair_solref: wp.array(dtype=wp.vec2), pair_solreffriction: wp.array(dtype=wp.vec2), @@ -251,7 +251,7 @@ def _box_box( geom_solimp[worldid], geom_friction[worldid], geom_margin[worldid], - geom_gap, + geom_gap[worldid], pair_dim, pair_solref, pair_solreffriction, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 71467bad..afed3625 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -715,7 +715,7 @@ def gjk_epa_sparse( geom_size: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d(dtype=wp.vec3), geom_margin: wp.array2d(dtype=float), - geom_gap: wp.array(dtype=float), + geom_gap: wp.array2d(dtype=float), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), mesh_vert: wp.array(dtype=wp.vec3), @@ -761,7 +761,7 @@ def gjk_epa_sparse( geom_solimp[worldid], geom_friction[worldid], geom_margin[worldid], - geom_gap, + geom_gap[worldid], pair_dim, pair_solref, pair_solreffriction, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index b7444661..37b0a6f4 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1109,7 +1109,7 @@ def contact_params( geom_solimp: wp.array(dtype=vec5), geom_friction: wp.array(dtype=wp.vec3), geom_margin: wp.array(dtype=float), - geom_gap: wp.array(dtype=float), + geom_gap: wp.array2d(dtype=float), pair_dim: wp.array(dtype=int), pair_solref: wp.array(dtype=wp.vec2), pair_solreffriction: wp.array(dtype=wp.vec2), @@ -1731,7 +1731,7 @@ def _primitive_narrowphase( geom_size: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d(dtype=wp.vec3), geom_margin: wp.array(dtype=float), - geom_gap: wp.array(dtype=float), + geom_gap: wp.array2d(dtype=float), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), pair_dim: wp.array(dtype=int), @@ -1776,7 +1776,7 @@ def _primitive_narrowphase( geom_solimp[worldid], geom_friction[worldid], geom_margin[worldid], - geom_gap, + geom_gap[worldid], pair_dim, pair_solref, pair_solreffriction, From 8ed2187126a0fff9a9e83d36f00357e026e7870e Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:26:44 +0200 Subject: [PATCH 36/74] site_pos --- 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 13891c73..ee949a8d 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -175,7 +175,7 @@ def _geom_local_to_global( def _site_local_to_global( # Model: site_bodyid: wp.array(dtype=int), - site_pos: wp.array(dtype=wp.vec3), + site_pos: wp.array2d(dtype=wp.vec3), site_quat: wp.array(dtype=wp.quat), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), @@ -188,7 +188,7 @@ def _site_local_to_global( bodyid = site_bodyid[siteid] xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] - site_xpos_out[worldid, siteid] = xpos + math.rot_vec_quat(site_pos[siteid], xquat) + site_xpos_out[worldid, siteid] = xpos + math.rot_vec_quat(site_pos[worldid, siteid], xquat) site_xmat_out[worldid, siteid] = math.quat_to_mat(math.mul_quat(xquat, site_quat[siteid])) @@ -1105,7 +1105,7 @@ def _cfrc_ext_equality( # Model: body_rootid: wp.array(dtype=int), site_bodyid: wp.array(dtype=int), - site_pos: wp.array(dtype=wp.vec3), + site_pos: wp.array2d(dtype=wp.vec3), eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), eq_objtype: wp.array(dtype=int), @@ -1168,7 +1168,7 @@ def _cfrc_ext_equality( else: offset = wp.vec3(eq_data_[3], eq_data_[4], eq_data_[5]) else: - offset = site_pos[obj1] + offset = site_pos[worldid, obj1] # transform point on body1: local -> global pos = xmat_in[worldid, bodyid1] @ offset + xpos_in[worldid, bodyid1] @@ -1190,7 +1190,7 @@ def _cfrc_ext_equality( else: offset = wp.vec3(eq_data_[0], eq_data_[1], eq_data_[2]) else: - offset = site_pos[obj2] + offset = site_pos[worldid,obj2] # transform point on body2: local -> global pos = xmat_in[worldid, bodyid2] @ offset + xpos_in[worldid, bodyid2] From cab4a48dd37cbd542df0785acabe4ef026441034 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:27:40 +0200 Subject: [PATCH 37/74] site_quat --- mujoco_warp/_src/constraint.py | 6 +++--- mujoco_warp/_src/sensor.py | 8 ++++---- mujoco_warp/_src/smooth.py | 4 ++-- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 7a15f700..09dc663c 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -492,7 +492,7 @@ def _efc_equality_weld( body_invweight0: wp.array3d(dtype=float), dof_bodyid: wp.array(dtype=int), site_bodyid: wp.array(dtype=int), - site_quat: wp.array(dtype=wp.quat), + site_quat: wp.array2d(dtype=wp.quat), eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), eq_objtype: wp.array(dtype=int), @@ -552,8 +552,8 @@ def _efc_equality_weld( pos1 = site_xpos_in[worldid, obj1id] pos2 = site_xpos_in[worldid, obj2id] - quat = math.mul_quat(xquat_in[worldid, body1id], site_quat[obj1id]) - quat1 = math.quat_inv(math.mul_quat(xquat_in[worldid, body2id], site_quat[obj2id])) + quat = math.mul_quat(xquat_in[worldid, body1id], site_quat[worldid, obj1id]) + quat1 = math.quat_inv(math.mul_quat(xquat_in[worldid, body2id], site_quat[worldid, obj2id])) else: body1id = obj1id diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 64afc290..bf86f67d 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -283,7 +283,7 @@ def _frame_quat( geom_bodyid: wp.array(dtype=int), geom_quat: wp.array2d(dtype=wp.quat), site_bodyid: wp.array(dtype=int), - site_quat: wp.array(dtype=wp.quat), + site_quat: wp.array2d(dtype=wp.quat), # Data in: xquat_in: wp.array2d(dtype=wp.quat), # In: @@ -308,10 +308,10 @@ def _frame_quat( return quat refquat = math.mul_quat(xquat_in[worldid, geom_bodyid[refid]], geom_quat[worldid,refid]) elif objtype == int(ObjType.SITE.value): - quat = math.mul_quat(xquat_in[worldid, site_bodyid[objid]], site_quat[objid]) + quat = math.mul_quat(xquat_in[worldid, site_bodyid[objid]], site_quat[worldid, objid]) if refid == -1: return quat - refquat = math.mul_quat(xquat_in[worldid, site_bodyid[refid]], site_quat[refid]) + refquat = math.mul_quat(xquat_in[worldid, site_bodyid[refid]], site_quat[worldid, refid]) # TODO(team): camera @@ -337,7 +337,7 @@ def _sensor_pos( body_iquat: wp.array2d(dtype=wp.quat), jnt_qposadr: wp.array(dtype=int), geom_bodyid: wp.array(dtype=int), - geom_quat: wp.array(dtype=wp.quat), + geom_quat: wp.array2d(dtype=wp.quat), site_bodyid: wp.array(dtype=int), site_quat: wp.array(dtype=wp.quat), cam_fovy: wp.array(dtype=float), diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index ee949a8d..67c194ae 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -176,7 +176,7 @@ def _site_local_to_global( # Model: site_bodyid: wp.array(dtype=int), site_pos: wp.array2d(dtype=wp.vec3), - site_quat: wp.array(dtype=wp.quat), + site_quat: wp.array2d(dtype=wp.quat), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), @@ -189,7 +189,7 @@ 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, siteid], xquat) - site_xmat_out[worldid, siteid] = math.quat_to_mat(math.mul_quat(xquat, site_quat[siteid])) + site_xmat_out[worldid, siteid] = math.quat_to_mat(math.mul_quat(xquat, site_quat[worldid,siteid])) @event_scope From 8a0d69e3d2be323af55d4cfb23071b458eda7bca Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:28:38 +0200 Subject: [PATCH 38/74] cam_pos --- 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 67c194ae..4201598c 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -437,7 +437,7 @@ def com_pos(m: Model, d: Data): def _cam_local_to_global( # Model: cam_bodyid: wp.array(dtype=int), - cam_pos: wp.array(dtype=wp.vec3), + cam_pos: wp.array2d(dtype=wp.vec3), cam_quat: wp.array(dtype=wp.quat), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), @@ -451,7 +451,7 @@ def _cam_local_to_global( bodyid = cam_bodyid[camid] xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] - cam_xpos_out[worldid, camid] = xpos + math.rot_vec_quat(cam_pos[camid], xquat) + cam_xpos_out[worldid, camid] = xpos + math.rot_vec_quat(cam_pos[worldid,camid], xquat) cam_xmat_out[worldid, camid] = math.quat_to_mat(math.mul_quat(xquat, cam_quat[camid])) From b777101b5a10c105c1ea63f5ea63c700bc6b9393 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:29:02 +0200 Subject: [PATCH 39/74] cam_quat --- 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 4201598c..1076b831 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -438,7 +438,7 @@ def _cam_local_to_global( # Model: cam_bodyid: wp.array(dtype=int), cam_pos: wp.array2d(dtype=wp.vec3), - cam_quat: wp.array(dtype=wp.quat), + cam_quat: wp.array2d(dtype=wp.quat), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), @@ -452,7 +452,7 @@ def _cam_local_to_global( xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] cam_xpos_out[worldid, camid] = xpos + math.rot_vec_quat(cam_pos[worldid,camid], xquat) - cam_xmat_out[worldid, camid] = math.quat_to_mat(math.mul_quat(xquat, cam_quat[camid])) + cam_xmat_out[worldid, camid] = math.quat_to_mat(math.mul_quat(xquat, cam_quat[worldid, camid])) @wp.kernel @@ -461,7 +461,7 @@ def _cam_fn( cam_mode: wp.array(dtype=int), cam_bodyid: wp.array(dtype=int), cam_targetbodyid: wp.array(dtype=int), - cam_poscom0: wp.array(dtype=wp.vec3), + cam_poscom0: wp.array2d(dtype=wp.vec3), cam_pos0: wp.array(dtype=wp.vec3), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), From bc662a7aa757aa25fe47fddd61be955e7cf8fe73 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:30:30 +0200 Subject: [PATCH 40/74] cam_pos0 and poscom0 --- 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 1076b831..de00ecb2 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -462,7 +462,7 @@ def _cam_fn( cam_bodyid: wp.array(dtype=int), cam_targetbodyid: wp.array(dtype=int), cam_poscom0: wp.array2d(dtype=wp.vec3), - cam_pos0: wp.array(dtype=wp.vec3), + cam_pos0: wp.array2d(dtype=wp.vec3), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), subtree_com_in: wp.array2d(dtype=wp.vec3), @@ -479,9 +479,9 @@ def _cam_fn( return elif cam_mode[camid] == wp.static(CamLightType.TRACK.value): body_xpos = xpos_in[worldid, cam_bodyid[camid]] - cam_xpos_out[worldid, camid] = body_xpos + cam_pos0[camid] + cam_xpos_out[worldid, camid] = body_xpos + cam_pos0[worldid,camid] elif cam_mode[camid] == wp.static(CamLightType.TRACKCOM.value): - cam_xpos_out[worldid, camid] = subtree_com_in[worldid, cam_bodyid[camid]] + cam_poscom0[camid] + cam_xpos_out[worldid, camid] = subtree_com_in[worldid, cam_bodyid[camid]] + cam_poscom0[worldid, camid] elif cam_mode[camid] == wp.static(CamLightType.TARGETBODY.value) or cam_mode[camid] == wp.static( CamLightType.TARGETBODYCOM.value ): From a624543d59b79599111a952ed223168293e13eee Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:31:26 +0200 Subject: [PATCH 41/74] light_pos --- 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 de00ecb2..45db51a4 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -506,7 +506,7 @@ def _cam_fn( def _light_local_to_global( # Model: light_bodyid: wp.array(dtype=int), - light_pos: wp.array(dtype=wp.vec3), + light_pos: wp.array2d(dtype=wp.vec3), light_dir: wp.array(dtype=wp.vec3), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), @@ -520,7 +520,7 @@ def _light_local_to_global( bodyid = light_bodyid[lightid] xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] - light_xpos_out[worldid, lightid] = xpos + math.rot_vec_quat(light_pos[lightid], xquat) + light_xpos_out[worldid, lightid] = xpos + math.rot_vec_quat(light_pos[worldid, lightid], xquat) light_xdir_out[worldid, lightid] = math.rot_vec_quat(light_dir[lightid], xquat) From 3d0a9e306d3b1744946d862a822f75bd7e4a5250 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:31:51 +0200 Subject: [PATCH 42/74] light_dir --- 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 45db51a4..e1393b7c 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -507,7 +507,7 @@ def _light_local_to_global( # Model: light_bodyid: wp.array(dtype=int), light_pos: wp.array2d(dtype=wp.vec3), - light_dir: wp.array(dtype=wp.vec3), + light_dir: wp.array2d(dtype=wp.vec3), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), xquat_in: wp.array2d(dtype=wp.quat), @@ -521,7 +521,7 @@ def _light_local_to_global( xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] light_xpos_out[worldid, lightid] = xpos + math.rot_vec_quat(light_pos[worldid, lightid], xquat) - light_xdir_out[worldid, lightid] = math.rot_vec_quat(light_dir[lightid], xquat) + light_xdir_out[worldid, lightid] = math.rot_vec_quat(light_dir[worldid,lightid], xquat) @wp.kernel From 12bb345ffcfe93b84b2580827a6abc87cb0a7da6 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:32:18 +0200 Subject: [PATCH 43/74] light_poscom0 --- 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 e1393b7c..e2967217 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -530,7 +530,7 @@ def _light_fn( light_mode: wp.array(dtype=int), light_bodyid: wp.array(dtype=int), light_targetbodyid: wp.array(dtype=int), - light_poscom0: wp.array(dtype=wp.vec3), + light_poscom0: wp.array2d(dtype=wp.vec3), light_pos0: wp.array(dtype=wp.vec3), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), @@ -551,7 +551,7 @@ def _light_fn( body_xpos = xpos_in[worldid, light_bodyid[lightid]] light_xpos_out[worldid, lightid] = body_xpos + light_pos0[lightid] elif light_mode[lightid] == wp.static(CamLightType.TRACKCOM.value): - light_xpos_out[worldid, lightid] = subtree_com_in[worldid, light_bodyid[lightid]] + light_poscom0[lightid] + light_xpos_out[worldid, lightid] = subtree_com_in[worldid, light_bodyid[lightid]] + light_poscom0[worldid, lightid] elif light_mode[lightid] == wp.static(CamLightType.TARGETBODY.value) or light_mode[lightid] == wp.static( CamLightType.TARGETBODYCOM.value ): From fb4b0cf5341d2c37969594d2798b27b37bfea864 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:32:44 +0200 Subject: [PATCH 44/74] light_pos0 --- 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 e2967217..8db03031 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -531,7 +531,7 @@ def _light_fn( light_bodyid: wp.array(dtype=int), light_targetbodyid: wp.array(dtype=int), light_poscom0: wp.array2d(dtype=wp.vec3), - light_pos0: wp.array(dtype=wp.vec3), + light_pos0: wp.array2d(dtype=wp.vec3), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), light_xpos_in: wp.array2d(dtype=wp.vec3), @@ -549,7 +549,7 @@ def _light_fn( return elif light_mode[lightid] == wp.static(CamLightType.TRACK.value): body_xpos = xpos_in[worldid, light_bodyid[lightid]] - light_xpos_out[worldid, lightid] = body_xpos + light_pos0[lightid] + light_xpos_out[worldid, lightid] = body_xpos + light_pos0[worldid, lightid] elif light_mode[lightid] == wp.static(CamLightType.TRACKCOM.value): light_xpos_out[worldid, lightid] = subtree_com_in[worldid, light_bodyid[lightid]] + light_poscom0[worldid, lightid] elif light_mode[lightid] == wp.static(CamLightType.TARGETBODY.value) or light_mode[lightid] == wp.static( From 912f90211776dc83f4518e180f96c62abc41e3ce Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:33:31 +0200 Subject: [PATCH 45/74] eq_solref --- mujoco_warp/_src/constraint.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 09dc663c..698cc6ab 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -105,7 +105,7 @@ def _efc_equality_connect( eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), eq_objtype: wp.array(dtype=int), - eq_solref: wp.array(dtype=wp.vec2), + eq_solref: wp.array2d(dtype=wp.vec2), eq_solimp: wp.array(dtype=vec5), eq_data: wp.array(dtype=vec11), eq_connect_adr: wp.array(dtype=int), @@ -198,7 +198,7 @@ def _efc_equality_connect( invweight = body_invweight0[worldid, body1id, 0] + body_invweight0[worldid, body2id, 0] pos_imp = wp.length(pos) - solref = eq_solref[i_eq] + solref = eq_solref[worldid, i_eq] solimp = eq_solimp[i_eq] for i in range(3): @@ -237,7 +237,7 @@ def _efc_equality_joint( dof_invweight0: wp.array2d(dtype=float), eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), - eq_solref: wp.array(dtype=wp.vec2), + eq_solref: wp.array2d(dtype=wp.vec2), eq_solimp: wp.array(dtype=vec5), eq_data: wp.array(dtype=vec11), eq_jnt_adr: wp.array(dtype=int), @@ -306,7 +306,7 @@ def _efc_equality_joint( pos, pos, invweight, - eq_solref[i_eq], + eq_solref[worldid, i_eq], eq_solimp[i_eq], 0.0, Jqvel, @@ -328,7 +328,7 @@ def _efc_equality_tendon( opt_timestep: float, eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), - eq_solref: wp.array(dtype=wp.vec2), + eq_solref: wp.array2d(dtype=wp.vec2), eq_solimp: wp.array(dtype=vec5), eq_data: wp.array(dtype=vec11), eq_ten_adr: wp.array(dtype=int), @@ -369,7 +369,7 @@ def _efc_equality_tendon( obj1id = eq_obj1id[eqid] obj2id = eq_obj2id[eqid] data = eq_data[eqid] - solref = eq_solref[eqid] + solref = eq_solref[worldid, eqid] solimp = eq_solimp[eqid] pos1 = ten_length_in[worldid, obj1id] - tendon_length0[obj1id] pos2 = ten_length_in[worldid, obj2id] - tendon_length0[obj2id] @@ -496,7 +496,7 @@ def _efc_equality_weld( eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), eq_objtype: wp.array(dtype=int), - eq_solref: wp.array(dtype=wp.vec2), + eq_solref: wp.array2d(dtype=wp.vec2), eq_solimp: wp.array(dtype=vec5), eq_data: wp.array(dtype=vec11), eq_wld_adr: wp.array(dtype=int), @@ -616,7 +616,7 @@ def _efc_equality_weld( pos_imp = wp.sqrt(wp.length_sq(cpos) + wp.length_sq(crot)) - solref = eq_solref[i_eq] + solref = eq_solref[worldid, i_eq] solimp = eq_solimp[i_eq] for i in range(3): From 45c93389cf1affc1c1382e6e0118b197c593a173 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:34:22 +0200 Subject: [PATCH 46/74] eq_solimp --- mujoco_warp/_src/constraint.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 698cc6ab..ac409efe 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -106,7 +106,7 @@ def _efc_equality_connect( eq_obj2id: wp.array(dtype=int), eq_objtype: wp.array(dtype=int), eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array(dtype=vec5), + eq_solimp: wp.array2d(dtype=vec5), eq_data: wp.array(dtype=vec11), eq_connect_adr: wp.array(dtype=int), # Data in: @@ -199,7 +199,7 @@ def _efc_equality_connect( pos_imp = wp.length(pos) solref = eq_solref[worldid, i_eq] - solimp = eq_solimp[i_eq] + solimp = eq_solimp[worldid, i_eq] for i in range(3): efcidi = efcid + i @@ -238,7 +238,7 @@ def _efc_equality_joint( eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array(dtype=vec5), + eq_solimp: wp.array2d(dtype=vec5), eq_data: wp.array(dtype=vec11), eq_jnt_adr: wp.array(dtype=int), # Data in: @@ -307,7 +307,7 @@ def _efc_equality_joint( pos, invweight, eq_solref[worldid, i_eq], - eq_solimp[i_eq], + eq_solimp[worldid, i_eq], 0.0, Jqvel, 0.0, @@ -329,7 +329,7 @@ def _efc_equality_tendon( eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array(dtype=vec5), + eq_solimp: wp.array2d(dtype=vec5), eq_data: wp.array(dtype=vec11), eq_ten_adr: wp.array(dtype=int), tendon_length0: wp.array(dtype=float), @@ -370,7 +370,7 @@ def _efc_equality_tendon( obj2id = eq_obj2id[eqid] data = eq_data[eqid] solref = eq_solref[worldid, eqid] - solimp = eq_solimp[eqid] + solimp = eq_solimp[worldid, eqid] pos1 = ten_length_in[worldid, obj1id] - tendon_length0[obj1id] pos2 = ten_length_in[worldid, obj2id] - tendon_length0[obj2id] jac1 = ten_J_in[worldid, obj1id] @@ -497,7 +497,7 @@ def _efc_equality_weld( eq_obj2id: wp.array(dtype=int), eq_objtype: wp.array(dtype=int), eq_solref: wp.array2d(dtype=wp.vec2), - eq_solimp: wp.array(dtype=vec5), + eq_solimp: wp.array2d(dtype=vec5), eq_data: wp.array(dtype=vec11), eq_wld_adr: wp.array(dtype=int), # Data in: @@ -617,7 +617,7 @@ def _efc_equality_weld( pos_imp = wp.sqrt(wp.length_sq(cpos) + wp.length_sq(crot)) solref = eq_solref[worldid, i_eq] - solimp = eq_solimp[i_eq] + solimp = eq_solimp[worldid, i_eq] for i in range(3): _update_efc_row( From fe3842d5d9af9bb5f24d0a9370adfc9a6ef8d2d2 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:36:09 +0200 Subject: [PATCH 47/74] actuator_dynprm --- mujoco_warp/_src/forward.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index a851a406..3f854b5c 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -141,7 +141,7 @@ def _next_activation( opt_timestep: float, actuator_dyntype: wp.array(dtype=int), actuator_actlimited: wp.array(dtype=bool), - actuator_dynprm: wp.array(dtype=vec10f), + actuator_dynprm: wp.array2d(dtype=vec10f), actuator_actrange: wp.array(dtype=wp.vec2), # Data in: act_in: wp.array2d(dtype=float), @@ -159,7 +159,7 @@ def _next_activation( # advance the actuation if actuator_dyntype[actid] == wp.static(DynType.FILTEREXACT.value): - dyn_prm = actuator_dynprm[actid] + dyn_prm = actuator_dynprm[worldid,actid] tau = wp.max(MJ_MINVAL, dyn_prm[0]) act += act_dot_scale_in * act_dot * tau * (1.0 - wp.exp(-opt_timestep / tau)) else: @@ -818,7 +818,7 @@ def _actuator_force( actuator_actnum: wp.array(dtype=int), actuator_ctrllimited: wp.array(dtype=bool), actuator_forcelimited: wp.array(dtype=bool), - actuator_dynprm: wp.array(dtype=vec10f), + actuator_dynprm: wp.array2d(dtype=vec10f), actuator_gainprm: wp.array(dtype=vec10f), actuator_biasprm: wp.array(dtype=vec10f), actuator_ctrlrange: wp.array(dtype=wp.vec2), @@ -848,7 +848,7 @@ def _actuator_force( if dyntype == int(DynType.INTEGRATOR.value): act_dot_out[worldid, actuator_actadr[uid]] = ctrl elif dyntype == int(DynType.FILTER.value) or dyntype == int(DynType.FILTEREXACT.value): - dynprm = actuator_dynprm[uid] + dynprm = actuator_dynprm[worldid, uid] actadr = actuator_actadr[uid] act = act_in[worldid, actadr] act_dot_out[worldid, actadr] = (ctrl - act) / wp.max(dynprm[0], MJ_MINVAL) From 467c1b68ac766187bea9a745629016ce926f980a Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:36:48 +0200 Subject: [PATCH 48/74] gainprm --- mujoco_warp/_src/forward.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index 3f854b5c..c2eaa385 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -481,7 +481,7 @@ def _implicit_actuator_bias_gain_vel( actuator_dyntype: wp.array(dtype=int), actuator_gaintype: wp.array(dtype=int), actuator_biastype: wp.array(dtype=int), - actuator_gainprm: wp.array(dtype=vec10f), + actuator_gainprm: wp.array2d(dtype=vec10f), actuator_biasprm: wp.array(dtype=vec10f), # Data in: act_in: wp.array2d(dtype=float), @@ -497,7 +497,7 @@ def _implicit_actuator_bias_gain_vel( bias_vel = 0.0 if actuator_gaintype[actid] == wp.static(GainType.AFFINE.value): - gain_vel = actuator_gainprm[actid][2] + gain_vel = actuator_gainprm[worldid,actid][2] else: gain_vel = 0.0 @@ -819,7 +819,7 @@ def _actuator_force( actuator_ctrllimited: wp.array(dtype=bool), actuator_forcelimited: wp.array(dtype=bool), actuator_dynprm: wp.array2d(dtype=vec10f), - actuator_gainprm: wp.array(dtype=vec10f), + actuator_gainprm: wp.array2d(dtype=vec10f), actuator_biasprm: wp.array(dtype=vec10f), actuator_ctrlrange: wp.array(dtype=wp.vec2), actuator_forcerange: wp.array(dtype=wp.vec2), @@ -867,7 +867,7 @@ def _actuator_force( # gain gaintype = actuator_gaintype[uid] - gainprm = actuator_gainprm[uid] + gainprm = actuator_gainprm[worldid, uid] gain = 0.0 if gaintype == int(GainType.FIXED.value): From 3eda7dd366bdd6be22fdf20b441a50114760ba86 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:37:25 +0200 Subject: [PATCH 49/74] biasprm --- mujoco_warp/_src/forward.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index c2eaa385..594ef91b 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -482,7 +482,7 @@ def _implicit_actuator_bias_gain_vel( actuator_gaintype: wp.array(dtype=int), actuator_biastype: wp.array(dtype=int), actuator_gainprm: wp.array2d(dtype=vec10f), - actuator_biasprm: wp.array(dtype=vec10f), + actuator_biasprm: wp.array2d(dtype=vec10f), # Data in: act_in: wp.array2d(dtype=float), ctrl_in: wp.array2d(dtype=float), @@ -492,7 +492,7 @@ def _implicit_actuator_bias_gain_vel( worldid, actid = wp.tid() if actuator_biastype[actid] == wp.static(BiasType.AFFINE.value): - bias_vel = actuator_biasprm[actid][2] + bias_vel = actuator_biasprm[worldid, actid][2] else: bias_vel = 0.0 @@ -820,7 +820,7 @@ def _actuator_force( actuator_forcelimited: wp.array(dtype=bool), actuator_dynprm: wp.array2d(dtype=vec10f), actuator_gainprm: wp.array2d(dtype=vec10f), - actuator_biasprm: wp.array(dtype=vec10f), + actuator_biasprm: wp.array2d(dtype=vec10f), actuator_ctrlrange: wp.array(dtype=wp.vec2), actuator_forcerange: wp.array(dtype=wp.vec2), # Data in: @@ -879,7 +879,7 @@ def _actuator_force( # bias biastype = actuator_biastype[uid] - biasprm = actuator_biasprm[uid] + biasprm = actuator_biasprm[worldid, uid] bias = 0.0 # BiasType.NONE if biastype == int(BiasType.AFFINE.value): From fd0f33f8a05a3d73535dd3f0b9b5c9cf952a29a7 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:37:56 +0200 Subject: [PATCH 50/74] actuator_ctrlrange --- mujoco_warp/_src/forward.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index 594ef91b..227495c4 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -821,7 +821,7 @@ def _actuator_force( actuator_dynprm: wp.array2d(dtype=vec10f), actuator_gainprm: wp.array2d(dtype=vec10f), actuator_biasprm: wp.array2d(dtype=vec10f), - actuator_ctrlrange: wp.array(dtype=wp.vec2), + actuator_ctrlrange: wp.array2d(dtype=wp.vec2), actuator_forcerange: wp.array(dtype=wp.vec2), # Data in: act_in: wp.array2d(dtype=float), @@ -839,7 +839,7 @@ def _actuator_force( ctrl = ctrl_in[worldid, uid] if actuator_ctrllimited[uid] and not dsbl_clampctrl: - ctrlrange = actuator_ctrlrange[uid] + ctrlrange = actuator_ctrlrange[worldid, uid] ctrl = wp.clamp(ctrl, ctrlrange[0], ctrlrange[1]) if na: From 65b34765a6f51e8bd32e9b221e1c8965f39617be Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:38:22 +0200 Subject: [PATCH 51/74] actuator_forcerange --- mujoco_warp/_src/forward.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index 227495c4..9ef2a724 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -822,7 +822,7 @@ def _actuator_force( actuator_gainprm: wp.array2d(dtype=vec10f), actuator_biasprm: wp.array2d(dtype=vec10f), actuator_ctrlrange: wp.array2d(dtype=wp.vec2), - actuator_forcerange: wp.array(dtype=wp.vec2), + actuator_forcerange: wp.array2d(dtype=wp.vec2), # Data in: act_in: wp.array2d(dtype=float), ctrl_in: wp.array2d(dtype=float), @@ -892,7 +892,7 @@ def _actuator_force( # TODO(team): tendon total force clamping if actuator_forcelimited[uid]: - forcerange = actuator_forcerange[uid] + forcerange = actuator_forcerange[worldid, uid] force = wp.clamp(force, forcerange[0], forcerange[1]) actuator_force_out[worldid, uid] = force From 4822742092b32e76ee31770df5c5e2926fa251a8 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:39:14 +0200 Subject: [PATCH 52/74] actrange and gear --- mujoco_warp/_src/forward.py | 4 ++-- mujoco_warp/_src/smooth.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index 9ef2a724..de346cd1 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -142,7 +142,7 @@ def _next_activation( actuator_dyntype: wp.array(dtype=int), actuator_actlimited: wp.array(dtype=bool), actuator_dynprm: wp.array2d(dtype=vec10f), - actuator_actrange: wp.array(dtype=wp.vec2), + actuator_actrange: wp.array2d(dtype=wp.vec2), # Data in: act_in: wp.array2d(dtype=float), act_dot_in: wp.array2d(dtype=float), @@ -167,7 +167,7 @@ def _next_activation( # clamp to actrange if limit and actuator_actlimited[actid]: - actrange = actuator_actrange[actid] + actrange = actuator_actrange[worldid,actid] act = wp.clamp(act, actrange[0], actrange[1]) act_out[worldid, actid] = act diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 8db03031..ec5da1ef 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -1445,7 +1445,7 @@ def _transmission( jnt_dofadr: wp.array(dtype=int), actuator_trntype: wp.array(dtype=int), actuator_trnid: wp.array(dtype=wp.vec2i), - actuator_gear: wp.array(dtype=wp.spatial_vector), + actuator_gear: wp.array2d(dtype=wp.spatial_vector), tendon_adr: wp.array(dtype=int), tendon_num: wp.array(dtype=int), wrap_objid: wp.array(dtype=int), @@ -1460,7 +1460,7 @@ def _transmission( ): worldid, actid = wp.tid() trntype = actuator_trntype[actid] - gear = actuator_gear[actid] + gear = actuator_gear[worldid, actid] if trntype == wp.static(TrnType.JOINT.value) or trntype == wp.static(TrnType.JOINTINPARENT.value): qpos = qpos_in[worldid] jntid = actuator_trnid[actid][0] From 48c1e777da5bb3f2765c3962156ee2ef94870db4 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:40:37 +0200 Subject: [PATCH 53/74] pair_solref --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 6 +++--- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 85b27fda..6ef82eb1 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -200,7 +200,7 @@ def _box_box( geom_margin: wp.array2d(dtype=float), geom_gap: wp.array2d(dtype=float), pair_dim: wp.array(dtype=int), - pair_solref: wp.array(dtype=wp.vec2), + pair_solref: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array(dtype=wp.vec2), pair_solimp: wp.array(dtype=vec5), pair_margin: wp.array(dtype=float), @@ -253,7 +253,7 @@ def _box_box( geom_margin[worldid], geom_gap[worldid], pair_dim, - pair_solref, + pair_solref[worldid], pair_solreffriction, pair_solimp, pair_margin, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index afed3625..f9d10cdc 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -720,7 +720,7 @@ def gjk_epa_sparse( mesh_vertnum: wp.array(dtype=int), mesh_vert: wp.array(dtype=wp.vec3), pair_dim: wp.array(dtype=int), - pair_solref: wp.array(dtype=wp.vec2), + pair_solref: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array(dtype=wp.vec2), pair_solimp: wp.array(dtype=vec5), pair_margin: wp.array(dtype=float), @@ -763,7 +763,7 @@ def gjk_epa_sparse( geom_margin[worldid], geom_gap[worldid], pair_dim, - pair_solref, + pair_solref[worldid], pair_solreffriction, pair_solimp, pair_margin, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 37b0a6f4..d2c3a31a 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1109,7 +1109,7 @@ def contact_params( geom_solimp: wp.array(dtype=vec5), geom_friction: wp.array(dtype=wp.vec3), geom_margin: wp.array(dtype=float), - geom_gap: wp.array2d(dtype=float), + geom_gap: wp.array(dtype=float), pair_dim: wp.array(dtype=int), pair_solref: wp.array(dtype=wp.vec2), pair_solreffriction: wp.array(dtype=wp.vec2), @@ -1735,7 +1735,7 @@ def _primitive_narrowphase( mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), pair_dim: wp.array(dtype=int), - pair_solref: wp.array(dtype=wp.vec2), + pair_solref: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array(dtype=wp.vec2), pair_solimp: wp.array(dtype=vec5), pair_margin: wp.array(dtype=float), @@ -1778,7 +1778,7 @@ def _primitive_narrowphase( geom_margin[worldid], geom_gap[worldid], pair_dim, - pair_solref, + pair_solref[worldid], pair_solreffriction, pair_solimp, pair_margin, From d12e94d67aa0f4605ef4549aeb1e1a9e2c9155e3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:42:17 +0200 Subject: [PATCH 54/74] pair_solreffriction --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 6ef82eb1..87439bfb 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -201,7 +201,7 @@ def _box_box( geom_gap: wp.array2d(dtype=float), pair_dim: wp.array(dtype=int), pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array(dtype=wp.vec2), + pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solimp: wp.array(dtype=vec5), pair_margin: wp.array(dtype=float), pair_gap: wp.array(dtype=float), @@ -254,7 +254,7 @@ def _box_box( geom_gap[worldid], pair_dim, pair_solref[worldid], - pair_solreffriction, + pair_solreffriction[worldid], pair_solimp, pair_margin, pair_gap, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index f9d10cdc..c6c63a2f 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -721,7 +721,7 @@ def gjk_epa_sparse( mesh_vert: wp.array(dtype=wp.vec3), pair_dim: wp.array(dtype=int), pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array(dtype=wp.vec2), + pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solimp: wp.array(dtype=vec5), pair_margin: wp.array(dtype=float), pair_gap: wp.array(dtype=float), @@ -764,7 +764,7 @@ def gjk_epa_sparse( geom_gap[worldid], pair_dim, pair_solref[worldid], - pair_solreffriction, + pair_solreffriction[worldid], pair_solimp, pair_margin, pair_gap, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index d2c3a31a..56576689 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1736,7 +1736,7 @@ def _primitive_narrowphase( mesh_vertnum: wp.array(dtype=int), pair_dim: wp.array(dtype=int), pair_solref: wp.array2d(dtype=wp.vec2), - pair_solreffriction: wp.array(dtype=wp.vec2), + pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solimp: wp.array(dtype=vec5), pair_margin: wp.array(dtype=float), pair_gap: wp.array(dtype=float), @@ -1779,7 +1779,7 @@ def _primitive_narrowphase( geom_gap[worldid], pair_dim, pair_solref[worldid], - pair_solreffriction, + pair_solreffriction[worldid], pair_solimp, pair_margin, pair_gap, From 584f263b8ff5a1c37a91ab520064a99c97c4a2ec Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:43:09 +0200 Subject: [PATCH 55/74] pair_solimp --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 87439bfb..f0decb14 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -202,7 +202,7 @@ def _box_box( pair_dim: wp.array(dtype=int), pair_solref: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array(dtype=vec5), + pair_solimp: wp.array2d(dtype=vec5), pair_margin: wp.array(dtype=float), pair_gap: wp.array(dtype=float), pair_friction: wp.array(dtype=vec5), @@ -255,7 +255,7 @@ def _box_box( pair_dim, pair_solref[worldid], pair_solreffriction[worldid], - pair_solimp, + pair_solimp[worldid], pair_margin, pair_gap, pair_friction, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index c6c63a2f..7db32370 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -722,7 +722,7 @@ def gjk_epa_sparse( pair_dim: wp.array(dtype=int), pair_solref: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array(dtype=vec5), + pair_solimp: wp.array2d(dtype=vec5), pair_margin: wp.array(dtype=float), pair_gap: wp.array(dtype=float), pair_friction: wp.array(dtype=vec5), @@ -765,7 +765,7 @@ def gjk_epa_sparse( pair_dim, pair_solref[worldid], pair_solreffriction[worldid], - pair_solimp, + pair_solimp[worldid], pair_margin, pair_gap, pair_friction, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 56576689..0a3c9ad5 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1737,7 +1737,7 @@ def _primitive_narrowphase( pair_dim: wp.array(dtype=int), pair_solref: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d(dtype=wp.vec2), - pair_solimp: wp.array(dtype=vec5), + pair_solimp: wp.array2d(dtype=vec5), pair_margin: wp.array(dtype=float), pair_gap: wp.array(dtype=float), pair_friction: wp.array(dtype=vec5), @@ -1780,7 +1780,7 @@ def _primitive_narrowphase( pair_dim, pair_solref[worldid], pair_solreffriction[worldid], - pair_solimp, + pair_solimp[worldid], pair_margin, pair_gap, pair_friction, From 0546318ccd0a7c05fac1af037da12aca1dcdc3e7 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:43:59 +0200 Subject: [PATCH 56/74] pair_margin --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index f0decb14..1bac0ed1 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -203,7 +203,7 @@ def _box_box( pair_solref: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array(dtype=float), + pair_margin: wp.array2d(dtype=float), pair_gap: wp.array(dtype=float), pair_friction: wp.array(dtype=vec5), # Data in: @@ -256,7 +256,7 @@ def _box_box( pair_solref[worldid], pair_solreffriction[worldid], pair_solimp[worldid], - pair_margin, + pair_margin[worldid], pair_gap, pair_friction, collision_pair_in, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 7db32370..45d58de9 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -723,7 +723,7 @@ def gjk_epa_sparse( pair_solref: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array(dtype=float), + pair_margin: wp.array2d(dtype=float), pair_gap: wp.array(dtype=float), pair_friction: wp.array(dtype=vec5), # Data in: @@ -766,7 +766,7 @@ def gjk_epa_sparse( pair_solref[worldid], pair_solreffriction[worldid], pair_solimp[worldid], - pair_margin, + pair_margin[worldid], pair_gap, pair_friction, collision_pair_in, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 0a3c9ad5..aacf17d2 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1738,7 +1738,7 @@ def _primitive_narrowphase( pair_solref: wp.array2d(dtype=wp.vec2), pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solimp: wp.array2d(dtype=vec5), - pair_margin: wp.array(dtype=float), + pair_margin: wp.array2d(dtype=float), pair_gap: wp.array(dtype=float), pair_friction: wp.array(dtype=vec5), # Data in: @@ -1781,7 +1781,7 @@ def _primitive_narrowphase( pair_solref[worldid], pair_solreffriction[worldid], pair_solimp[worldid], - pair_margin, + pair_margin[worldid], pair_gap, pair_friction, collision_pair_in, From 77f6921ecc121b06ee0ae768b925875d3ba6eb37 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:44:45 +0200 Subject: [PATCH 57/74] pair_gap --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 1bac0ed1..f026c59d 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -204,7 +204,7 @@ def _box_box( pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solimp: wp.array2d(dtype=vec5), pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array(dtype=float), + pair_gap: wp.array2d(dtype=float), pair_friction: wp.array(dtype=vec5), # Data in: nconmax_in: int, @@ -257,7 +257,7 @@ def _box_box( pair_solreffriction[worldid], pair_solimp[worldid], pair_margin[worldid], - pair_gap, + pair_gap[worldid], pair_friction, collision_pair_in, collision_pairid_in, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 45d58de9..cd4b4809 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -724,7 +724,7 @@ def gjk_epa_sparse( pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solimp: wp.array2d(dtype=vec5), pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array(dtype=float), + pair_gap: wp.array2d(dtype=float), pair_friction: wp.array(dtype=vec5), # Data in: nconmax_in: int, @@ -767,7 +767,7 @@ def gjk_epa_sparse( pair_solreffriction[worldid], pair_solimp[worldid], pair_margin[worldid], - pair_gap, + pair_gap[worldid], pair_friction, collision_pair_in, collision_pairid_in, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index aacf17d2..13b2bdfc 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1739,7 +1739,7 @@ def _primitive_narrowphase( pair_solreffriction: wp.array2d(dtype=wp.vec2), pair_solimp: wp.array2d(dtype=vec5), pair_margin: wp.array2d(dtype=float), - pair_gap: wp.array(dtype=float), + pair_gap: wp.array2d(dtype=float), pair_friction: wp.array(dtype=vec5), # Data in: nconmax_in: int, @@ -1782,7 +1782,7 @@ def _primitive_narrowphase( pair_solreffriction[worldid], pair_solimp[worldid], pair_margin[worldid], - pair_gap, + pair_gap[worldid], pair_friction, collision_pair_in, collision_pairid_in, From 5c03632c0b21c1be79dd394bec5e3f049a888cda Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:45:49 +0200 Subject: [PATCH 58/74] pair_friction --- mujoco_warp/_src/collision_box.py | 4 ++-- mujoco_warp/_src/collision_convex.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index f026c59d..5b91f999 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -205,7 +205,7 @@ def _box_box( pair_solimp: wp.array2d(dtype=vec5), pair_margin: wp.array2d(dtype=float), pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array(dtype=vec5), + pair_friction: wp.array2d(dtype=vec5), # Data in: nconmax_in: int, geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -258,7 +258,7 @@ def _box_box( pair_solimp[worldid], pair_margin[worldid], pair_gap[worldid], - pair_friction, + pair_friction[worldid], collision_pair_in, collision_pairid_in, tid, diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index cd4b4809..41e1e804 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -725,7 +725,7 @@ def gjk_epa_sparse( pair_solimp: wp.array2d(dtype=vec5), pair_margin: wp.array2d(dtype=float), pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array(dtype=vec5), + pair_friction: wp.array2d(dtype=vec5), # Data in: nconmax_in: int, geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -768,7 +768,7 @@ def gjk_epa_sparse( pair_solimp[worldid], pair_margin[worldid], pair_gap[worldid], - pair_friction, + pair_friction[worldid], collision_pair_in, collision_pairid_in, tid, diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 13b2bdfc..15967272 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1740,7 +1740,7 @@ def _primitive_narrowphase( pair_solimp: wp.array2d(dtype=vec5), pair_margin: wp.array2d(dtype=float), pair_gap: wp.array2d(dtype=float), - pair_friction: wp.array(dtype=vec5), + pair_friction: wp.array2d(dtype=vec5), # Data in: nconmax_in: int, geom_xpos_in: wp.array2d(dtype=wp.vec3), @@ -1783,7 +1783,7 @@ def _primitive_narrowphase( pair_solimp[worldid], pair_margin[worldid], pair_gap[worldid], - pair_friction, + pair_friction[worldid], collision_pair_in, collision_pairid_in, tid, From 95d1fe9a723c32550f112ef0f37d73953ae320af Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:47:42 +0200 Subject: [PATCH 59/74] tendon --- mujoco_warp/_src/constraint.py | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index ac409efe..d41cd741 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -332,8 +332,8 @@ def _efc_equality_tendon( eq_solimp: wp.array2d(dtype=vec5), eq_data: wp.array(dtype=vec11), eq_ten_adr: wp.array(dtype=int), - tendon_length0: wp.array(dtype=float), - tendon_invweight0: wp.array(dtype=float), + tendon_length0: wp.array2d(dtype=float), + tendon_invweight0: wp.array2d(dtype=float), # Data in: ne_connect_in: wp.array(dtype=int), ne_weld_in: wp.array(dtype=int), @@ -371,13 +371,13 @@ def _efc_equality_tendon( data = eq_data[eqid] solref = eq_solref[worldid, eqid] solimp = eq_solimp[worldid, eqid] - pos1 = ten_length_in[worldid, obj1id] - tendon_length0[obj1id] - pos2 = ten_length_in[worldid, obj2id] - tendon_length0[obj2id] + pos1 = ten_length_in[worldid, obj1id] - tendon_length0[worldid, obj1id] + pos2 = ten_length_in[worldid, obj2id] - tendon_length0[worldid, obj2id] jac1 = ten_J_in[worldid, obj1id] jac2 = ten_J_in[worldid, obj2id] if obj2id > -1: - invweight = tendon_invweight0[obj1id] + tendon_invweight0[obj2id] + invweight = tendon_invweight0[worldid, obj1id] + tendon_invweight0[worldid, obj2id] dif = pos2 dif2 = dif * dif @@ -387,7 +387,7 @@ def _efc_equality_tendon( pos = pos1 - (data[0] + data[1] * dif + data[2] * dif2 + data[3] * dif3 + data[4] * dif4) deriv = data[1] + 2.0 * data[2] * dif + 3.0 * data[3] * dif2 + 4.0 * data[4] * dif3 else: - invweight = tendon_invweight0[obj1id] + invweight = tendon_invweight0[worldid, obj1id] pos = pos1 - data[0] deriv = 0.0 @@ -827,11 +827,11 @@ def _efc_limit_tendon( tendon_adr: wp.array(dtype=int), tendon_num: wp.array(dtype=int), tendon_limited_adr: wp.array(dtype=int), - tendon_solref_lim: wp.array(dtype=wp.vec2), - tendon_solimp_lim: wp.array(dtype=vec5), - tendon_range: wp.array(dtype=wp.vec2), - tendon_margin: wp.array(dtype=float), - tendon_invweight0: wp.array(dtype=float), + tendon_solref_lim: wp.array2d(dtype=wp.vec2), + tendon_solimp_lim: wp.array2d(dtype=vec5), + tendon_range: wp.array2d(dtype=wp.vec2), + tendon_margin: wp.array2d(dtype=float), + tendon_invweight0: wp.array2d(dtype=float), wrap_objid: wp.array(dtype=int), wrap_type: wp.array(dtype=int), # Data in: @@ -855,10 +855,10 @@ def _efc_limit_tendon( worldid, tenlimitedid = wp.tid() tenid = tendon_limited_adr[tenlimitedid] - tenrange = tendon_range[tenid] + tenrange = tendon_range[worldid, tenid] length = ten_length_in[worldid, tenid] dist_min, dist_max = length - tenrange[0], tenrange[1] - length - tenmargin = tendon_margin[tenid] + tenmargin = tendon_margin[worldid,tenid] pos = wp.min(dist_min, dist_max) - tenmargin active = pos < 0 @@ -890,9 +890,9 @@ def _efc_limit_tendon( efcid, pos, pos, - tendon_invweight0[tenid], - tendon_solref_lim[tenid], - tendon_solimp_lim[tenid], + tendon_invweight0[worldid,tenid], + tendon_solref_lim[worldid, tenid], + tendon_solimp_lim[worldid, tenid], tenmargin, Jqvel, 0.0, From d6f8ef98bb2836db164d920d1c3abf7fd728eae6 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:48:10 +0200 Subject: [PATCH 60/74] wrap prm --- 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 ec5da1ef..b0f88040 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -1870,7 +1870,7 @@ def _joint_tendon( jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), wrap_objid: wp.array(dtype=int), - wrap_prm: wp.array(dtype=float), + wrap_prm: wp.array2d(dtype=float), tendon_jnt_adr: wp.array(dtype=int), wrap_jnt_adr: wp.array(dtype=int), # Data in: @@ -1885,7 +1885,7 @@ def _joint_tendon( wrap_jnt_adr_ = wrap_jnt_adr[wrapid] wrap_objid_ = wrap_objid[wrap_jnt_adr_] - prm = wrap_prm[wrap_jnt_adr_] + prm = wrap_prm[worldid, wrap_jnt_adr_] # add to length L = prm * qpos_in[worldid, jnt_qposadr[wrap_objid_]] From 2f2bc49daaade13183aaee23fb899bb667bd64e0 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 17:50:05 +0200 Subject: [PATCH 61/74] sensor_cutoff --- mujoco_warp/_src/sensor.py | 62 +++++++++++++++++++------------------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index bf86f67d..18413354 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -339,7 +339,7 @@ def _sensor_pos( geom_bodyid: wp.array(dtype=int), geom_quat: wp.array2d(dtype=wp.quat), site_bodyid: wp.array(dtype=int), - site_quat: wp.array(dtype=wp.quat), + site_quat: wp.array2d(dtype=wp.quat), cam_fovy: wp.array(dtype=float), cam_resolution: wp.array(dtype=wp.vec2i), cam_sensorsize: wp.array(dtype=wp.vec2), @@ -350,7 +350,7 @@ def _sensor_pos( sensor_objid: wp.array(dtype=int), sensor_refid: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), + sensor_cutoff: wp.array2d(dtype=float), sensor_pos_adr: wp.array(dtype=int), # Data in: time_in: wp.array(dtype=float), @@ -383,19 +383,19 @@ def _sensor_pos( vec2 = _cam_projection( cam_fovy, cam_resolution, cam_sensorsize, cam_intrinsic, site_xpos_in, cam_xpos_in, cam_xmat_in, worldid, objid, refid ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 2, vec2, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 2, vec2, out) elif sensortype == int(SensorType.JOINTPOS.value): val = _joint_pos(jnt_qposadr, qpos_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) elif sensortype == int(SensorType.TENDONPOS.value): val = _tendon_pos(ten_length_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) elif sensortype == int(SensorType.ACTUATORPOS.value): val = _actuator_pos(actuator_length_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) elif sensortype == int(SensorType.BALLQUAT.value): quat = _ball_quat(jnt_qposadr, qpos_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 4, quat, out) elif sensortype == int(SensorType.FRAMEPOS.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -413,7 +413,7 @@ def _sensor_pos( objtype, refid, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif ( sensortype == int(SensorType.FRAMEXAXIS.value) or sensortype == int(SensorType.FRAMEYAXIS.value) @@ -428,18 +428,18 @@ def _sensor_pos( elif sensortype == int(SensorType.FRAMEZAXIS.value): axis = 2 vec3 = _frame_axis(ximat_in, xmat_in, geom_xmat_in, site_xmat_in, worldid, objid, objtype, refid, axis) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.FRAMEQUAT.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] quat = _frame_quat(body_iquat, geom_bodyid, geom_quat, site_bodyid, site_quat, xquat_in, worldid, objid, objtype, refid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 4, quat, out) elif sensortype == int(SensorType.SUBTREECOM.value): vec3 = _subtree_com(subtree_com_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.CLOCK.value): val = _clock(time_in, worldid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) @event_scope @@ -811,7 +811,7 @@ def _sensor_vel( sensor_reftype: wp.array(dtype=int), sensor_refid: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), + sensor_cutoff: wp.array2d(dtype=float), sensor_vel_adr: wp.array(dtype=int), # Data in: qvel_in: wp.array2d(dtype=float), @@ -842,22 +842,22 @@ def _sensor_vel( if sensortype == int(SensorType.VELOCIMETER.value): vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.GYRO.value): vec3 = _gyro(site_bodyid, site_xmat_in, cvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.JOINTVEL.value): val = _joint_vel(jnt_dofadr, qvel_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) elif sensortype == int(SensorType.TENDONVEL.value): val = _tendon_vel(ten_velocity_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) elif sensortype == int(SensorType.ACTUATORVEL.value): val = _actuator_vel(actuator_velocity_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) elif sensortype == int(SensorType.BALLANGVEL.value): vec3 = _ball_ang_vel(jnt_dofadr, qvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.FRAMELINVEL.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -885,7 +885,7 @@ def _sensor_vel( refid, reftype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_linvel, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, frame_linvel, out) elif sensortype == int(SensorType.FRAMEANGVEL.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -913,13 +913,13 @@ def _sensor_vel( refid, reftype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_angvel, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, frame_angvel, out) elif sensortype == int(SensorType.SUBTREELINVEL.value): vec3 = _subtree_linvel(subtree_linvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.SUBTREEANGMOM.value): vec3 = _subtree_angmom(subtree_angmom_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) @event_scope @@ -1144,7 +1144,7 @@ def _sensor_acc( sensor_objtype: wp.array(dtype=int), sensor_objid: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), + sensor_cutoff: wp.array2d(dtype=float), sensor_acc_adr: wp.array(dtype=int), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), @@ -1171,19 +1171,19 @@ def _sensor_acc( vec3 = _accelerometer( body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, cacc_in, worldid, objid ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.FORCE.value): vec3 = _force(site_bodyid, site_xmat_in, cfrc_int_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.TORQUE.value): vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cfrc_int_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.ACTUATORFRC.value): val = _actuator_force(actuator_force_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) elif sensortype == int(SensorType.JOINTACTFRC.value): val = _joint_actuator_force(jnt_dofadr, qfrc_actuator_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) elif sensortype == int(SensorType.FRAMELINACC.value): objtype = sensor_objtype[sensorid] vec3 = _framelinacc( @@ -1201,7 +1201,7 @@ def _sensor_acc( objid, objtype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) elif sensortype == int(SensorType.FRAMEANGACC.value): objtype = sensor_objtype[sensorid] vec3 = _frameangacc( @@ -1212,7 +1212,7 @@ def _sensor_acc( objid, objtype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) @event_scope From a63d4e15e19a755b2060b02844bbc0ca3ec0da0b Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 18:07:33 +0200 Subject: [PATCH 62/74] fixes --- mujoco_warp/_src/collision_primitive.py | 6 +++--- mujoco_warp/_src/smooth.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 15967272..21a08608 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1730,7 +1730,7 @@ def _primitive_narrowphase( geom_solimp: wp.array2d(dtype=vec5), geom_size: wp.array2d(dtype=wp.vec3), geom_friction: wp.array2d(dtype=wp.vec3), - geom_margin: wp.array(dtype=float), + geom_margin: wp.array2d(dtype=float), geom_gap: wp.array2d(dtype=float), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), @@ -1767,6 +1767,8 @@ def _primitive_narrowphase( if tid >= ncollision_in[0]: return + + worldid = collision_worldid_in[tid] geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, @@ -1791,8 +1793,6 @@ def _primitive_narrowphase( g1 = geoms[0] g2 = geoms[1] - worldid = collision_worldid_in[tid] - geom1 = _geom( geom_dataid, geom_size[worldid], diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index b0f88040..d6886b50 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -280,12 +280,12 @@ def _subtree_com_acc( @wp.kernel def _subtree_div( # Model: - subtree_mass: wp.array(dtype=float), + subtree_mass: wp.array2d(dtype=float), # Data out: subtree_com_out: wp.array2d(dtype=wp.vec3), ): worldid, bodyid = wp.tid() - subtree_com_out[worldid, bodyid] /= subtree_mass[bodyid] + subtree_com_out[worldid, bodyid] /= subtree_mass[worldid,bodyid] @wp.kernel From 77a962ccf42badabe06abf363cb5bd3d4b5795f9 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 18:23:40 +0200 Subject: [PATCH 63/74] make kernel analyzer happy --- mujoco_warp/_src/collision_box.py | 27 ++--- mujoco_warp/_src/collision_convex.py | 27 ++--- mujoco_warp/_src/collision_primitive.py | 98 ++++++++--------- mujoco_warp/_src/types.py | 136 ++++++++++++------------ 4 files changed, 146 insertions(+), 142 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 5b91f999..0ea0c2f0 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -245,23 +245,24 @@ def _box_box( geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, - geom_priority[worldid], - geom_solmix[worldid], - geom_solref[worldid], - geom_solimp[worldid], - geom_friction[worldid], - geom_margin[worldid], - geom_gap[worldid], + geom_priority, + geom_solmix, + geom_solref, + geom_solimp, + geom_friction, + geom_margin, + geom_gap, pair_dim, - pair_solref[worldid], - pair_solreffriction[worldid], - pair_solimp[worldid], - pair_margin[worldid], - pair_gap[worldid], - pair_friction[worldid], + pair_solref, + pair_solreffriction, + pair_solimp, + pair_margin, + pair_gap, + pair_friction, collision_pair_in, collision_pairid_in, tid, + worldid ) # transformations diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 41e1e804..0d8ae744 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -755,23 +755,24 @@ def gjk_epa_sparse( worldid = collision_worldid_in[tid] geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, - geom_priority[worldid], - geom_solmix[worldid], - geom_solref[worldid], - geom_solimp[worldid], - geom_friction[worldid], - geom_margin[worldid], - geom_gap[worldid], + geom_priority, + geom_solmix, + geom_solref, + geom_solimp, + geom_friction, + geom_margin, + geom_gap, pair_dim, - pair_solref[worldid], - pair_solreffriction[worldid], - pair_solimp[worldid], - pair_margin[worldid], - pair_gap[worldid], - pair_friction[worldid], + pair_solref, + pair_solreffriction, + pair_solimp, + pair_margin, + pair_gap, + pair_friction, collision_pair_in, collision_pairid_in, tid, + worldid ) g1 = geoms[0] diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 21a08608..f276d971 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -42,7 +42,7 @@ class Geom: def _geom( # Model: geom_dataid: wp.array(dtype=int), - geom_size: wp.array(dtype=wp.vec3), + geom_size: wp.array2d(dtype=wp.vec3), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), # Data in: @@ -56,7 +56,7 @@ def _geom( geom.pos = geom_xpos_in[worldid, gid] rot = geom_xmat_in[worldid, gid] geom.rot = rot - geom.size = geom_size[gid] + geom.size = geom_size[worldid,gid] geom.normal = wp.vec3(rot[0, 2], rot[1, 2], rot[2, 2]) # plane dataid = geom_dataid[gid] @@ -1103,46 +1103,47 @@ def plane_cylinder( def contact_params( # Model: geom_condim: wp.array(dtype=int), - geom_priority: wp.array(dtype=int), - geom_solmix: wp.array(dtype=float), - geom_solref: wp.array(dtype=wp.vec2), - geom_solimp: wp.array(dtype=vec5), - geom_friction: wp.array(dtype=wp.vec3), - geom_margin: wp.array(dtype=float), - geom_gap: wp.array(dtype=float), + geom_priority: wp.array2d(dtype=int), + geom_solmix: wp.array2d(dtype=float), + geom_solref: wp.array2d(dtype=wp.vec2), + geom_solimp: wp.array2d(dtype=vec5), + geom_friction: wp.array2d(dtype=wp.vec3), + geom_margin: wp.array2d(dtype=float), + geom_gap: wp.array2d(dtype=float), pair_dim: wp.array(dtype=int), - pair_solref: wp.array(dtype=wp.vec2), - pair_solreffriction: wp.array(dtype=wp.vec2), - pair_solimp: wp.array(dtype=vec5), - pair_margin: wp.array(dtype=float), - pair_gap: wp.array(dtype=float), - pair_friction: wp.array(dtype=vec5), + pair_solref: wp.array2d(dtype=wp.vec2), + pair_solreffriction: wp.array2d(dtype=wp.vec2), + pair_solimp: wp.array2d(dtype=vec5), + pair_margin: wp.array2d(dtype=float), + pair_gap: wp.array2d(dtype=float), + pair_friction: wp.array2d(dtype=vec5), # Data in: collision_pair_in: wp.array(dtype=wp.vec2i), collision_pairid_in: wp.array(dtype=int), # In: cid: int, + worldid: int, ): geoms = collision_pair_in[cid] pairid = collision_pairid_in[cid] if pairid > -1: - margin = pair_margin[pairid] - gap = pair_gap[pairid] - condim = pair_dim[pairid] - friction = pair_friction[pairid] - solref = pair_solref[pairid] - solreffriction = pair_solreffriction[pairid] - solimp = pair_solimp[pairid] + margin = pair_margin[worldid, pairid] + gap = pair_gap[worldid, pairid] + condim = pair_dim[worldid, pairid] + friction = pair_friction[worldid, pairid] + solref = pair_solref[worldid, pairid] + solreffriction = pair_solreffriction[worldid, pairid] + solimp = pair_solimp[worldid, pairid] else: g1 = geoms[0] g2 = geoms[1] - p1 = geom_priority[g1] - p2 = geom_priority[g2] + p1 = geom_priority[worldid, g1] + p2 = geom_priority[worldid, g2] - solmix1 = geom_solmix[g1] - solmix2 = geom_solmix[g2] + solmix1 = geom_solmix[worldid, g1] + solmix2 = geom_solmix[worldid, g2] mix = solmix1 / (solmix1 + solmix2) mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix) @@ -1150,14 +1151,14 @@ def contact_params( mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix) mix = wp.where(p1 == p2, mix, wp.where(p1 > p2, 1.0, 0.0)) - margin = wp.max(geom_margin[g1], geom_margin[g2]) - gap = wp.max(geom_gap[g1], geom_gap[g2]) + margin = wp.max(geom_margin[worldid, g1], geom_margin[worldid, g2]) + gap = wp.max(geom_gap[worldid, g1], geom_gap[worldid, g2]) condim1 = geom_condim[g1] condim2 = geom_condim[g2] condim = wp.where(p1 == p2, wp.max(condim1, condim2), wp.where(p1 > p2, condim1, condim2)) - max_geom_friction = wp.max(geom_friction[g1], geom_friction[g2]) + max_geom_friction = wp.max(geom_friction[worldid, g1], geom_friction[worldid, g2]) friction = vec5( max_geom_friction[0], max_geom_friction[0], @@ -1166,14 +1167,14 @@ def contact_params( max_geom_friction[2], ) - if geom_solref[g1].x > 0.0 and geom_solref[g2].x > 0.0: - solref = mix * geom_solref[g1] + (1.0 - mix) * geom_solref[g2] + if geom_solref[worldid, g1].x > 0.0 and geom_solref[worldid, g2].x > 0.0: + solref = mix * geom_solref[worldid, g1] + (1.0 - mix) * geom_solref[worldid, g2] else: - solref = wp.min(geom_solref[g1], geom_solref[g2]) + solref = wp.min(geom_solref[worldid, g1], geom_solref[worldid, g2]) solreffriction = wp.vec2(0.0, 0.0) - solimp = mix * geom_solimp[g1] + (1.0 - mix) * geom_solimp[g2] + solimp = mix * geom_solimp[worldid, g1] + (1.0 - mix) * geom_solimp[worldid, g2] return geoms, margin, gap, condim, friction, solref, solreffriction, solimp @@ -1772,30 +1773,31 @@ def _primitive_narrowphase( geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( geom_condim, - geom_priority[worldid], - geom_solmix[worldid], - geom_solref[worldid], - geom_solimp[worldid], - geom_friction[worldid], - geom_margin[worldid], - geom_gap[worldid], + geom_priority, + geom_solmix, + geom_solref, + geom_solimp, + geom_friction, + geom_margin, + geom_gap, pair_dim, - pair_solref[worldid], - pair_solreffriction[worldid], - pair_solimp[worldid], - pair_margin[worldid], - pair_gap[worldid], - pair_friction[worldid], + pair_solref, + pair_solreffriction, + pair_solimp, + pair_margin, + pair_gap, + pair_friction, collision_pair_in, collision_pairid_in, tid, + worldid ) g1 = geoms[0] g2 = geoms[1] geom1 = _geom( geom_dataid, - geom_size[worldid], + geom_size, mesh_vertadr, mesh_vertnum, geom_xpos_in, @@ -1805,7 +1807,7 @@ def _primitive_narrowphase( ) geom2 = _geom( geom_dataid, - geom_size[worldid], + geom_size, mesh_vertadr, mesh_vertnum, geom_xpos_in, diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 58aebd0f..727fa31e 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -800,8 +800,8 @@ class Model: nlsp: int # warp only opt: Option stat: Statistic - qpos0: wp.array(dtype=float) - qpos_spring: wp.array(dtype=float) + qpos0: wp.array2d(dtype=float) + qpos_spring: wp.array2d(dtype=float) qM_fullm_i: wp.array(dtype=int) # warp only qM_fullm_j: wp.array(dtype=int) # warp only qM_mulm_i: wp.array(dtype=int) # warp only @@ -824,15 +824,15 @@ class Model: body_dofadr: wp.array(dtype=int) body_geomnum: wp.array(dtype=int) body_geomadr: wp.array(dtype=int) - body_pos: wp.array(dtype=wp.vec3) - body_quat: wp.array(dtype=wp.quat) - body_ipos: wp.array(dtype=wp.vec3) - body_iquat: wp.array(dtype=wp.quat) - body_mass: wp.array(dtype=float) - body_subtreemass: wp.array(dtype=float) - subtree_mass: wp.array(dtype=float) - body_inertia: wp.array(dtype=wp.vec3) - body_invweight0: wp.array2d(dtype=float) + body_pos: wp.array2d(dtype=wp.vec3) + body_quat: wp.array2d(dtype=wp.quat) + body_ipos: wp.array2d(dtype=wp.vec3) + body_iquat: wp.array2d(dtype=wp.quat) + body_mass: wp.array2d(dtype=float) + body_subtreemass: wp.array2d(dtype=float) + subtree_mass: wp.array2d(dtype=float) + body_inertia: wp.array2d(dtype=wp.vec3) + body_invweight0: wp.array3d(dtype=float) body_contype: wp.array(dtype=int) body_conaffinity: wp.array(dtype=int) body_gravcomp: wp.array(dtype=float) @@ -842,14 +842,14 @@ class Model: jnt_bodyid: wp.array(dtype=int) jnt_limited: wp.array(dtype=int) jnt_actfrclimited: wp.array(dtype=bool) - jnt_solref: wp.array(dtype=wp.vec2) - jnt_solimp: wp.array(dtype=vec5) - jnt_pos: wp.array(dtype=wp.vec3) + jnt_solref: wp.array2d(dtype=wp.vec2) + jnt_solimp: wp.array2d(dtype=vec5) + jnt_pos: wp.array2d(dtype=wp.vec3) jnt_axis: wp.array(dtype=wp.vec3) - jnt_stiffness: wp.array(dtype=float) - jnt_range: wp.array2d(dtype=float) - jnt_actfrcrange: wp.array(dtype=wp.vec2) - jnt_margin: wp.array(dtype=float) + jnt_stiffness: wp.array2d(dtype=float) + jnt_range: wp.array3d(dtype=float) + jnt_actfrcrange: wp.array2d(dtype=wp.vec2) + jnt_margin: wp.array2d(dtype=float) jnt_limited_slide_hinge_adr: wp.array(dtype=int) # warp only jnt_limited_ball_adr: wp.array(dtype=int) # warp only jnt_actgravcomp: wp.array(dtype=int) @@ -857,12 +857,12 @@ class Model: dof_jntid: wp.array(dtype=int) dof_parentid: wp.array(dtype=int) dof_Madr: wp.array(dtype=int) - dof_armature: wp.array(dtype=float) - dof_damping: wp.array(dtype=float) - dof_invweight0: wp.array(dtype=float) - dof_frictionloss: wp.array(dtype=float) - dof_solimp: wp.array(dtype=vec5) - dof_solref: wp.array(dtype=wp.vec2) + dof_armature: wp.array2d(dtype=float) + dof_damping: wp.array2d(dtype=float) + dof_invweight0: wp.array2d(dtype=float) + dof_frictionloss: wp.array2d(dtype=float) + dof_solimp: wp.array2d(dtype=vec5) + dof_solref: wp.array2d(dtype=wp.vec2) dof_tri_row: wp.array(dtype=int) # warp only dof_tri_col: wp.array(dtype=int) # warp only geom_type: wp.array(dtype=int) @@ -871,28 +871,28 @@ class Model: geom_condim: wp.array(dtype=int) geom_bodyid: wp.array(dtype=int) geom_dataid: wp.array(dtype=int) - geom_priority: wp.array(dtype=int) - geom_solmix: wp.array(dtype=float) - geom_solref: wp.array(dtype=wp.vec2) - geom_solimp: wp.array(dtype=vec5) - geom_size: wp.array(dtype=wp.vec3) + geom_priority: wp.array2d(dtype=int) + geom_solmix: wp.array2d(dtype=float) + geom_solref: wp.array2d(dtype=wp.vec2) + geom_solimp: wp.array2d(dtype=vec5) + geom_size: wp.array2d(dtype=wp.vec3) geom_aabb: wp.array2d(dtype=wp.vec3) - geom_rbound: wp.array(dtype=float) - geom_pos: wp.array(dtype=wp.vec3) - geom_quat: wp.array(dtype=wp.quat) - geom_friction: wp.array(dtype=wp.vec3) - geom_margin: wp.array(dtype=float) - geom_gap: wp.array(dtype=float) + geom_rbound: wp.array2d(dtype=float) + geom_pos: wp.array2d(dtype=wp.vec3) + geom_quat: wp.array2d(dtype=wp.quat) + geom_friction: wp.array2d(dtype=wp.vec3) + geom_margin: wp.array2d(dtype=float) + geom_gap: wp.array2d(dtype=float) site_bodyid: wp.array(dtype=int) - site_pos: wp.array(dtype=wp.vec3) - site_quat: wp.array(dtype=wp.quat) + site_pos: wp.array2d(dtype=wp.vec3) + site_quat: wp.array2d(dtype=wp.quat) cam_mode: wp.array(dtype=int) cam_bodyid: wp.array(dtype=int) cam_targetbodyid: wp.array(dtype=int) - cam_pos: wp.array(dtype=wp.vec3) - cam_quat: wp.array(dtype=wp.quat) - cam_poscom0: wp.array(dtype=wp.vec3) - cam_pos0: wp.array(dtype=wp.vec3) + cam_pos: wp.array2d(dtype=wp.vec3) + cam_quat: wp.array2d(dtype=wp.quat) + cam_poscom0: wp.array2d(dtype=wp.vec3) + cam_pos0: wp.array2d(dtype=wp.vec3) cam_fovy: wp.array(dtype=float) cam_resolution: wp.array(dtype=wp.vec2i) cam_sensorsize: wp.array(dtype=wp.vec2) @@ -900,10 +900,10 @@ class Model: light_mode: wp.array(dtype=int) light_bodyid: wp.array(dtype=int) light_targetbodyid: wp.array(dtype=int) - light_pos: wp.array(dtype=wp.vec3) - light_dir: wp.array(dtype=wp.vec3) - light_poscom0: wp.array(dtype=wp.vec3) - light_pos0: wp.array(dtype=wp.vec3) + light_pos: wp.array2d(dtype=wp.vec3) + light_dir: wp.array2d(dtype=wp.vec3) + light_poscom0: wp.array2d(dtype=wp.vec3) + light_pos0: wp.array2d(dtype=wp.vec3) mesh_vertadr: wp.array(dtype=int) mesh_vertnum: wp.array(dtype=int) mesh_vert: wp.array(dtype=wp.vec3) @@ -912,8 +912,8 @@ class Model: eq_obj2id: wp.array(dtype=int) eq_objtype: wp.array(dtype=int) eq_active0: wp.array(dtype=bool) - eq_solref: wp.array(dtype=wp.vec2) - eq_solimp: wp.array(dtype=vec5) + eq_solref: wp.array2d(dtype=wp.vec2) + eq_solimp: wp.array2d(dtype=vec5) eq_data: wp.array(dtype=vec11) eq_connect_adr: wp.array(dtype=int) eq_wld_adr: wp.array(dtype=int) @@ -932,38 +932,38 @@ class Model: actuator_ctrllimited: wp.array(dtype=bool) actuator_forcelimited: wp.array(dtype=bool) actuator_actlimited: wp.array(dtype=bool) - actuator_dynprm: wp.array(dtype=vec10f) - actuator_gainprm: wp.array(dtype=vec10f) - actuator_biasprm: wp.array(dtype=vec10f) - actuator_ctrlrange: wp.array(dtype=wp.vec2) - actuator_forcerange: wp.array(dtype=wp.vec2) - actuator_actrange: wp.array(dtype=wp.vec2) - actuator_gear: wp.array(dtype=wp.spatial_vector) + actuator_dynprm: wp.array2d(dtype=vec10f) + actuator_gainprm: wp.array2d(dtype=vec10f) + actuator_biasprm: wp.array2d(dtype=vec10f) + actuator_ctrlrange: wp.array2d(dtype=wp.vec2) + actuator_forcerange: wp.array2d(dtype=wp.vec2) + actuator_actrange: wp.array2d(dtype=wp.vec2) + actuator_gear: wp.array2d(dtype=wp.spatial_vector) nxn_geom_pair: wp.array(dtype=wp.vec2i) # warp only nxn_pairid: wp.array(dtype=int) # warp only pair_dim: wp.array(dtype=int) pair_geom1: wp.array(dtype=int) pair_geom2: wp.array(dtype=int) - pair_solref: wp.array(dtype=wp.vec2) - pair_solreffriction: wp.array(dtype=wp.vec2) - pair_solimp: wp.array(dtype=vec5) - pair_margin: wp.array(dtype=float) - pair_gap: wp.array(dtype=float) - pair_friction: wp.array(dtype=vec5) + pair_solref: wp.array2d(dtype=wp.vec2) + pair_solreffriction: wp.array2d(dtype=wp.vec2) + pair_solimp: wp.array2d(dtype=vec5) + pair_margin: wp.array2d(dtype=float) + pair_gap: wp.array2d(dtype=float) + pair_friction: wp.array2d(dtype=vec5) exclude_signature: wp.array(dtype=int) condim_max: int # warp only tendon_adr: wp.array(dtype=int) tendon_num: wp.array(dtype=int) tendon_limited: wp.array(dtype=int) tendon_limited_adr: wp.array(dtype=int) - tendon_solref_lim: wp.array(dtype=wp.vec2) - tendon_solimp_lim: wp.array(dtype=vec5) - tendon_range: wp.array(dtype=wp.vec2) - tendon_margin: wp.array(dtype=float) - tendon_length0: wp.array(dtype=float) - tendon_invweight0: wp.array(dtype=float) + tendon_solref_lim: wp.array2d(dtype=wp.vec2) + tendon_solimp_lim: wp.array2d(dtype=vec5) + tendon_range: wp.array2d(dtype=wp.vec2) + tendon_margin: wp.array2d(dtype=float) + tendon_length0: wp.array2d(dtype=float) + tendon_invweight0: wp.array2d(dtype=float) wrap_objid: wp.array(dtype=int) - wrap_prm: wp.array(dtype=float) + wrap_prm: wp.array2d(dtype=float) wrap_type: wp.array(dtype=int) tendon_jnt_adr: wp.array(dtype=int) # warp only tendon_site_adr: wp.array(dtype=int) # warp only @@ -981,7 +981,7 @@ class Model: sensor_refid: wp.array(dtype=int) sensor_dim: wp.array(dtype=int) sensor_adr: wp.array(dtype=int) - sensor_cutoff: wp.array(dtype=float) + sensor_cutoff: wp.array2d(dtype=float) sensor_pos_adr: wp.array(dtype=int) # warp only sensor_vel_adr: wp.array(dtype=int) # warp only sensor_acc_adr: wp.array(dtype=int) # warp only From e04ed2626c0c9ab47dccd638f7607a547721bf06 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 18:33:44 +0200 Subject: [PATCH 64/74] making kernel analyzer even more happy --- contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py b/contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py index 498377e4..87cb15b6 100644 --- a/contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py +++ b/contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py @@ -136,8 +136,8 @@ def test_all_issues( @kernel def test_no_issues( # Model: - qpos0: wp.array(dtype=float), - geom_pos: wp.array(dtype=wp.vec3), + qpos0: wp.array2d(dtype=float), + geom_pos: wp.array2d(dtype=wp.vec3), # Data in: qpos_in: wp.array2d(dtype=float), qvel_in: wp.array2d(dtype=float), diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index f276d971..7707243b 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1130,7 +1130,7 @@ def contact_params( if pairid > -1: margin = pair_margin[worldid, pairid] gap = pair_gap[worldid, pairid] - condim = pair_dim[worldid, pairid] + condim = pair_dim[pairid] friction = pair_friction[worldid, pairid] solref = pair_solref[worldid, pairid] solreffriction = pair_solreffriction[worldid, pairid] From 14fba1bd0a768f4c829f343177f7a3e4e7f18e95 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Thu, 8 May 2025 18:38:55 +0200 Subject: [PATCH 65/74] more fixes --- mujoco_warp/_src/sensor.py | 64 ++++++++++++++++++++------------------ 1 file changed, 33 insertions(+), 31 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 18413354..278f090a 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -35,15 +35,16 @@ def _write_scalar( # Model: sensor_datatype: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), + sensor_cutoff: wp.array2d(dtype=float), # In: sensorid: int, sensor: Any, + worldid: int, # Out: out: wp.array(dtype=float), ): adr = sensor_adr[sensorid] - cutoff = sensor_cutoff[sensorid] + cutoff = sensor_cutoff[worldid,sensorid] if cutoff > 0.0: datatype = sensor_datatype[sensorid] @@ -60,16 +61,17 @@ def _write_vector( # Model: sensor_datatype: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array(dtype=float), + sensor_cutoff: wp.array2d(dtype=float), # In: sensorid: int, sensordim: int, sensor: Any, + worldid: int, # Out: out: wp.array(dtype=float), ): adr = sensor_adr[sensorid] - cutoff = sensor_cutoff[sensorid] + cutoff = sensor_cutoff[worldid, sensorid] if cutoff > 0.0: datatype = sensor_datatype[sensorid] @@ -383,19 +385,19 @@ def _sensor_pos( vec2 = _cam_projection( cam_fovy, cam_resolution, cam_sensorsize, cam_intrinsic, site_xpos_in, cam_xpos_in, cam_xmat_in, worldid, objid, refid ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 2, vec2, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 2, vec2, worldid, out) elif sensortype == int(SensorType.JOINTPOS.value): val = _joint_pos(jnt_qposadr, qpos_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) elif sensortype == int(SensorType.TENDONPOS.value): val = _tendon_pos(ten_length_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) elif sensortype == int(SensorType.ACTUATORPOS.value): val = _actuator_pos(actuator_length_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) elif sensortype == int(SensorType.BALLQUAT.value): quat = _ball_quat(jnt_qposadr, qpos_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 4, quat, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, worldid, out) elif sensortype == int(SensorType.FRAMEPOS.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -413,7 +415,7 @@ def _sensor_pos( objtype, refid, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif ( sensortype == int(SensorType.FRAMEXAXIS.value) or sensortype == int(SensorType.FRAMEYAXIS.value) @@ -428,18 +430,18 @@ def _sensor_pos( elif sensortype == int(SensorType.FRAMEZAXIS.value): axis = 2 vec3 = _frame_axis(ximat_in, xmat_in, geom_xmat_in, site_xmat_in, worldid, objid, objtype, refid, axis) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.FRAMEQUAT.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] quat = _frame_quat(body_iquat, geom_bodyid, geom_quat, site_bodyid, site_quat, xquat_in, worldid, objid, objtype, refid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 4, quat, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, worldid, out) elif sensortype == int(SensorType.SUBTREECOM.value): vec3 = _subtree_com(subtree_com_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.CLOCK.value): val = _clock(time_in, worldid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) @event_scope @@ -842,22 +844,22 @@ def _sensor_vel( if sensortype == int(SensorType.VELOCIMETER.value): vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.GYRO.value): vec3 = _gyro(site_bodyid, site_xmat_in, cvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.JOINTVEL.value): val = _joint_vel(jnt_dofadr, qvel_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) elif sensortype == int(SensorType.TENDONVEL.value): val = _tendon_vel(ten_velocity_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) elif sensortype == int(SensorType.ACTUATORVEL.value): val = _actuator_vel(actuator_velocity_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) elif sensortype == int(SensorType.BALLANGVEL.value): vec3 = _ball_ang_vel(jnt_dofadr, qvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.FRAMELINVEL.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -885,7 +887,7 @@ def _sensor_vel( refid, reftype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, frame_linvel, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_linvel, worldid, out) elif sensortype == int(SensorType.FRAMEANGVEL.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -913,13 +915,13 @@ def _sensor_vel( refid, reftype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, frame_angvel, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, frame_angvel, worldid, out) elif sensortype == int(SensorType.SUBTREELINVEL.value): vec3 = _subtree_linvel(subtree_linvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.SUBTREEANGMOM.value): vec3 = _subtree_angmom(subtree_angmom_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, worldid, out) @event_scope @@ -1171,19 +1173,19 @@ def _sensor_acc( vec3 = _accelerometer( body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, cacc_in, worldid, objid ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.FORCE.value): vec3 = _force(site_bodyid, site_xmat_in, cfrc_int_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.TORQUE.value): vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cfrc_int_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.ACTUATORFRC.value): val = _actuator_force(actuator_force_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) elif sensortype == int(SensorType.JOINTACTFRC.value): val = _joint_actuator_force(jnt_dofadr, qfrc_actuator_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, val, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) elif sensortype == int(SensorType.FRAMELINACC.value): objtype = sensor_objtype[sensorid] vec3 = _framelinacc( @@ -1201,7 +1203,7 @@ def _sensor_acc( objid, objtype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.FRAMEANGACC.value): objtype = sensor_objtype[sensorid] vec3 = _frameangacc( @@ -1212,7 +1214,7 @@ def _sensor_acc( objid, objtype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) @event_scope From 24b91ada2741203ea09978e4d55ac196aefc7b0f Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Fri, 9 May 2025 10:14:44 +0200 Subject: [PATCH 66/74] small fix --- mujoco_warp/_src/smooth.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 876cb68b..aed9eb41 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -195,8 +195,8 @@ def _site_local_to_global( @wp.kernel def _mocap( # Model: - body_ipos: wp.array(dtype=wp.vec3), - body_iquat: wp.array(dtype=wp.quat), + body_ipos: wp.array2d(dtype=wp.vec3), + body_iquat: wp.array2d(dtype=wp.quat), mocap_bodyid: wp.array(dtype=int), # Data in: mocap_pos_in: wp.array2d(dtype=wp.vec3), @@ -215,8 +215,8 @@ def _mocap( xpos_out[worldid, bodyid] = xpos xquat_out[worldid, bodyid] = mocap_quat xmat_out[worldid, bodyid] = math.quat_to_mat(mocap_quat) - xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[bodyid], mocap_quat) - ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(mocap_quat, body_iquat[bodyid])) + xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid,bodyid], mocap_quat) + ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(mocap_quat, body_iquat[worldid, bodyid])) @event_scope From beacb46d6ba1a54d94d69237bc0d18718ea0a883 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Fri, 9 May 2025 10:17:11 +0200 Subject: [PATCH 67/74] fix formatting --- mujoco_warp/_src/collision_box.py | 2 +- mujoco_warp/_src/collision_convex.py | 2 +- mujoco_warp/_src/collision_driver.py | 4 ++-- mujoco_warp/_src/collision_primitive.py | 6 +++--- mujoco_warp/_src/constraint.py | 10 +++++----- mujoco_warp/_src/forward.py | 8 ++++---- mujoco_warp/_src/passive.py | 4 ++-- mujoco_warp/_src/sensor.py | 4 ++-- mujoco_warp/_src/smooth.py | 24 ++++++++++++------------ 9 files changed, 32 insertions(+), 32 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 0ea0c2f0..5c4f2038 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -262,7 +262,7 @@ def _box_box( collision_pair_in, collision_pairid_in, tid, - worldid + worldid, ) # transformations diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 0d8ae744..39b9461d 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -772,7 +772,7 @@ def gjk_epa_sparse( collision_pair_in, collision_pairid_in, tid, - worldid + worldid, ) g1 = geoms[0] diff --git a/mujoco_warp/_src/collision_driver.py b/mujoco_warp/_src/collision_driver.py index 26f16d96..be3a4732 100644 --- a/mujoco_warp/_src/collision_driver.py +++ b/mujoco_warp/_src/collision_driver.py @@ -138,13 +138,13 @@ def _sap_project( worldid, geomid = wp.tid() xpos = geom_xpos_in[worldid, geomid] - rbound = geom_rbound[worldid,geomid] + rbound = geom_rbound[worldid, geomid] if rbound == 0.0: # geom is a plane rbound = MJ_MAXVAL - radius = rbound + geom_margin[worldid,geomid] + radius = rbound + geom_margin[worldid, geomid] center = wp.dot(direction_in, xpos) sap_projection_lower_out[worldid, geomid] = center - radius diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 7707243b..8a47dbfc 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -56,7 +56,7 @@ def _geom( geom.pos = geom_xpos_in[worldid, gid] rot = geom_xmat_in[worldid, gid] geom.rot = rot - geom.size = geom_size[worldid,gid] + geom.size = geom_size[worldid, gid] geom.normal = wp.vec3(rot[0, 2], rot[1, 2], rot[2, 2]) # plane dataid = geom_dataid[gid] @@ -1768,7 +1768,7 @@ def _primitive_narrowphase( if tid >= ncollision_in[0]: return - + worldid = collision_worldid_in[tid] geoms, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params( @@ -1790,7 +1790,7 @@ def _primitive_narrowphase( collision_pair_in, collision_pairid_in, tid, - worldid + worldid, ) g1 = geoms[0] g2 = geoms[1] diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index d41cd741..4341f505 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -287,7 +287,7 @@ def _efc_equality_joint( rhs = data[0] + dif * (data[1] + dif * (data[2] + dif * (data[3] + dif * data[4]))) deriv_2 = data[1] + dif * (2.0 * data[2] + dif * (3.0 * data[3] + dif * 4.0 * data[4])) - pos = qpos_in[worldid, qposadr1] - qpos0[worldid,qposadr1] - rhs + pos = qpos_in[worldid, qposadr1] - qpos0[worldid, qposadr1] - rhs Jqvel = qvel_in[worldid, dofadr1] - qvel_in[worldid, dofadr2] * deriv_2 invweight = dof_invweight0[worldid, dofadr1] + dof_invweight0[worldid, dofadr2] @@ -697,7 +697,7 @@ def _efc_limit_slide_hinge( ): worldid, jntlimitedid = wp.tid() jntid = jnt_limited_slide_hinge_adr[jntlimitedid] - jntrange = jnt_range[worldid,jntid] + jntrange = jnt_range[worldid, jntid] qpos = qpos_in[worldid, jnt_qposadr[jntid]] jntmargin = jnt_margin[worldid, jntid] @@ -724,7 +724,7 @@ def _efc_limit_slide_hinge( pos, dof_invweight0[worldid, dofadr], jnt_solref[worldid, jntid], - jnt_solimp[worldid,jntid], + jnt_solimp[worldid, jntid], jntmargin, Jqvel, 0.0, @@ -858,7 +858,7 @@ def _efc_limit_tendon( tenrange = tendon_range[worldid, tenid] length = ten_length_in[worldid, tenid] dist_min, dist_max = length - tenrange[0], tenrange[1] - length - tenmargin = tendon_margin[worldid,tenid] + tenmargin = tendon_margin[worldid, tenid] pos = wp.min(dist_min, dist_max) - tenmargin active = pos < 0 @@ -890,7 +890,7 @@ def _efc_limit_tendon( efcid, pos, pos, - tendon_invweight0[worldid,tenid], + tendon_invweight0[worldid, tenid], tendon_solref_lim[worldid, tenid], tendon_solimp_lim[worldid, tenid], tenmargin, diff --git a/mujoco_warp/_src/forward.py b/mujoco_warp/_src/forward.py index de346cd1..c3c9cb7e 100644 --- a/mujoco_warp/_src/forward.py +++ b/mujoco_warp/_src/forward.py @@ -159,7 +159,7 @@ def _next_activation( # advance the actuation if actuator_dyntype[actid] == wp.static(DynType.FILTEREXACT.value): - dyn_prm = actuator_dynprm[worldid,actid] + dyn_prm = actuator_dynprm[worldid, actid] tau = wp.max(MJ_MINVAL, dyn_prm[0]) act += act_dot_scale_in * act_dot * tau * (1.0 - wp.exp(-opt_timestep / tau)) else: @@ -167,7 +167,7 @@ def _next_activation( # clamp to actrange if limit and actuator_actlimited[actid]: - actrange = actuator_actrange[worldid,actid] + actrange = actuator_actrange[worldid, actid] act = wp.clamp(act, actrange[0], actrange[1]) act_out[worldid, actid] = act @@ -278,7 +278,7 @@ def _euler_damp_qfrc_sparse( worldid, tid = wp.tid() adr = dof_Madr[tid] - qM_integration_out[worldid, 0, adr] += opt_timestep * dof_damping[worldid,tid] + qM_integration_out[worldid, 0, adr] += opt_timestep * dof_damping[worldid, tid] qfrc_integration_out[worldid, tid] = qfrc_smooth_in[worldid, tid] + qfrc_constraint_in[worldid, tid] @@ -497,7 +497,7 @@ def _implicit_actuator_bias_gain_vel( bias_vel = 0.0 if actuator_gaintype[actid] == wp.static(GainType.AFFINE.value): - gain_vel = actuator_gainprm[worldid,actid][2] + gain_vel = actuator_gainprm[worldid, actid][2] else: gain_vel = 0.0 diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 52f84204..51c7b325 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -38,7 +38,7 @@ def _spring_passive( qfrc_spring_out: wp.array2d(dtype=float), ): worldid, jntid = wp.tid() - stiffness = jnt_stiffness[worldid,jntid] + stiffness = jnt_stiffness[worldid, jntid] dofid = jnt_dofadr[jntid] if stiffness == 0.0: @@ -90,7 +90,7 @@ def _spring_passive( qfrc_spring_out[worldid, dofid + 1] = -stiffness * dif[1] qfrc_spring_out[worldid, dofid + 2] = -stiffness * dif[2] else: # mjJNT_SLIDE, mjJNT_HINGE - fdif = qpos_in[worldid, qposid] - qpos_spring[worldid,qposid] + fdif = qpos_in[worldid, qposid] - qpos_spring[worldid, qposid] qfrc_spring_out[worldid, dofid] = -stiffness * fdif diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 278f090a..49eaa996 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -44,7 +44,7 @@ def _write_scalar( out: wp.array(dtype=float), ): adr = sensor_adr[sensorid] - cutoff = sensor_cutoff[worldid,sensorid] + cutoff = sensor_cutoff[worldid, sensorid] if cutoff > 0.0: datatype = sensor_datatype[sensorid] @@ -308,7 +308,7 @@ def _frame_quat( quat = math.mul_quat(xquat_in[worldid, geom_bodyid[objid]], geom_quat[worldid, objid]) if refid == -1: return quat - refquat = math.mul_quat(xquat_in[worldid, geom_bodyid[refid]], geom_quat[worldid,refid]) + refquat = math.mul_quat(xquat_in[worldid, geom_bodyid[refid]], geom_quat[worldid, refid]) elif objtype == int(ObjType.SITE.value): quat = math.mul_quat(xquat_in[worldid, site_bodyid[objid]], site_quat[worldid, objid]) if refid == -1: diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index aed9eb41..87b2e96d 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -147,7 +147,7 @@ def _kinematics_level( 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, bodyid], xquat) - ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[worldid,bodyid])) + ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(xquat, body_iquat[worldid, bodyid])) @wp.kernel @@ -167,7 +167,7 @@ def _geom_local_to_global( bodyid = geom_bodyid[geomid] xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] - geom_xpos_out[worldid, geomid] = xpos + math.rot_vec_quat(geom_pos[worldid,geomid], xquat) + geom_xpos_out[worldid, geomid] = xpos + math.rot_vec_quat(geom_pos[worldid, geomid], xquat) geom_xmat_out[worldid, geomid] = math.quat_to_mat(math.mul_quat(xquat, geom_quat[worldid, geomid])) @@ -189,7 +189,7 @@ 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, siteid], xquat) - site_xmat_out[worldid, siteid] = math.quat_to_mat(math.mul_quat(xquat, site_quat[worldid,siteid])) + site_xmat_out[worldid, siteid] = math.quat_to_mat(math.mul_quat(xquat, site_quat[worldid, siteid])) @wp.kernel @@ -215,7 +215,7 @@ def _mocap( xpos_out[worldid, bodyid] = xpos xquat_out[worldid, bodyid] = mocap_quat xmat_out[worldid, bodyid] = math.quat_to_mat(mocap_quat) - xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid,bodyid], mocap_quat) + xipos_out[worldid, bodyid] = xpos + math.rot_vec_quat(body_ipos[worldid, bodyid], mocap_quat) ximat_out[worldid, bodyid] = math.quat_to_mat(math.mul_quat(mocap_quat, body_iquat[worldid, bodyid])) @@ -320,7 +320,7 @@ def _subtree_div( subtree_com_out: wp.array2d(dtype=wp.vec3), ): worldid, bodyid = wp.tid() - subtree_com_out[worldid, bodyid] /= subtree_mass[worldid,bodyid] + subtree_com_out[worldid, bodyid] /= subtree_mass[worldid, bodyid] @wp.kernel @@ -338,7 +338,7 @@ def _cinert( ): worldid, bodyid = wp.tid() mat = ximat_in[worldid, bodyid] - inert = body_inertia[worldid,bodyid] + inert = body_inertia[worldid, bodyid] mass = body_mass[worldid, bodyid] dif = xipos_in[worldid, bodyid] - subtree_com_in[worldid, body_rootid[bodyid]] # express inertia in com-based frame (mju_inertCom) @@ -486,7 +486,7 @@ def _cam_local_to_global( bodyid = cam_bodyid[camid] xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] - cam_xpos_out[worldid, camid] = xpos + math.rot_vec_quat(cam_pos[worldid,camid], xquat) + cam_xpos_out[worldid, camid] = xpos + math.rot_vec_quat(cam_pos[worldid, camid], xquat) cam_xmat_out[worldid, camid] = math.quat_to_mat(math.mul_quat(xquat, cam_quat[worldid, camid])) @@ -514,7 +514,7 @@ def _cam_fn( return elif cam_mode[camid] == wp.static(CamLightType.TRACK.value): body_xpos = xpos_in[worldid, cam_bodyid[camid]] - cam_xpos_out[worldid, camid] = body_xpos + cam_pos0[worldid,camid] + cam_xpos_out[worldid, camid] = body_xpos + cam_pos0[worldid, camid] elif cam_mode[camid] == wp.static(CamLightType.TRACKCOM.value): cam_xpos_out[worldid, camid] = subtree_com_in[worldid, cam_bodyid[camid]] + cam_poscom0[worldid, camid] elif cam_mode[camid] == wp.static(CamLightType.TARGETBODY.value) or cam_mode[camid] == wp.static( @@ -556,7 +556,7 @@ def _light_local_to_global( xpos = xpos_in[worldid, bodyid] xquat = xquat_in[worldid, bodyid] light_xpos_out[worldid, lightid] = xpos + math.rot_vec_quat(light_pos[worldid, lightid], xquat) - light_xdir_out[worldid, lightid] = math.rot_vec_quat(light_dir[worldid,lightid], xquat) + light_xdir_out[worldid, lightid] = math.rot_vec_quat(light_dir[worldid, lightid], xquat) @wp.kernel @@ -682,7 +682,7 @@ def _qM_sparse( bodyid = dof_bodyid[dofid] # init M(i,i) with armature inertia - qM_out[worldid, 0, madr_ij] = dof_armature[worldid,dofid] + qM_out[worldid, 0, madr_ij] = dof_armature[worldid, dofid] # precompute buf = crb_body_i * cdof_i buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid]) @@ -1225,7 +1225,7 @@ def _cfrc_ext_equality( else: offset = wp.vec3(eq_data_[0], eq_data_[1], eq_data_[2]) else: - offset = site_pos[worldid,obj2] + offset = site_pos[worldid, obj2] # transform point on body2: local -> global pos = xmat_in[worldid, bodyid2] @ offset + xpos_in[worldid, bodyid2] @@ -1797,7 +1797,7 @@ def _linear_momentum( if bodyid: pid = body_parentid[bodyid] wp.atomic_add(subtree_linvel_out[worldid], pid, subtree_linvel_in[worldid, bodyid]) - subtree_linvel_out[worldid, bodyid] /= wp.max(MJ_MINVAL, body_subtreemass[worldid,bodyid]) + subtree_linvel_out[worldid, bodyid] /= wp.max(MJ_MINVAL, body_subtreemass[worldid, bodyid]) @wp.kernel From dd57b09c2b1e817ebfb035b0ddcf26d8601992e3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Fri, 9 May 2025 14:25:02 +0200 Subject: [PATCH 68/74] fixes for sensor_test --- 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 49eaa996..6f6aedca 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -915,13 +915,13 @@ def _sensor_vel( refid, reftype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, frame_angvel, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_angvel, worldid, out) elif sensortype == int(SensorType.SUBTREELINVEL.value): vec3 = _subtree_linvel(subtree_linvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) elif sensortype == int(SensorType.SUBTREEANGMOM.value): vec3 = _subtree_angmom(subtree_angmom_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff[worldid], sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) @event_scope From 90b2cd783a942656b4f216647b924812a0ba2bae Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Mon, 12 May 2025 10:14:00 +0200 Subject: [PATCH 69/74] fixes after merge --- mujoco_warp/_src/sensor.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 95b20114..660a836a 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -300,7 +300,7 @@ def _frame_quat( site_bodyid: wp.array(dtype=int), site_quat: wp.array2d(dtype=wp.quat), cam_bodyid: wp.array(dtype=int), - cam_quat: wp.array(dtype=wp.quat), + cam_quat: wp.array2d(dtype=wp.quat), # Data in: xquat_in: wp.array2d(dtype=wp.quat), # In: @@ -327,15 +327,15 @@ def _frame_quat( return quat if reftype == int(ObjType.BODY.value): - refquat = math.mul_quat(xquat_in[worldid, refid], body_iquat[refid]) + refquat = math.mul_quat(xquat_in[worldid, refid], body_iquat[worldid, refid]) elif reftype == int(ObjType.XBODY.value): refquat = xquat_in[worldid, refid] elif reftype == int(ObjType.GEOM.value): - refquat = math.mul_quat(xquat_in[worldid, geom_bodyid[refid]], geom_quat[refid]) + refquat = math.mul_quat(xquat_in[worldid, geom_bodyid[refid]], geom_quat[worldid, refid]) elif reftype == int(ObjType.SITE.value): - refquat = math.mul_quat(xquat_in[worldid, site_bodyid[refid]], site_quat[refid]) + refquat = math.mul_quat(xquat_in[worldid, site_bodyid[refid]], site_quat[worldid, refid]) elif reftype == int(ObjType.CAMERA.value): - refquat = math.mul_quat(xquat_in[worldid, cam_bodyid[refid]], cam_quat[refid]) + refquat = math.mul_quat(xquat_in[worldid, cam_bodyid[refid]], cam_quat[worldid, refid]) else: # UNKNOWN refquat = wp.quat(1.0, 0.0, 0.0, 0.0) @@ -362,7 +362,7 @@ def _sensor_pos( site_bodyid: wp.array(dtype=int), site_quat: wp.array2d(dtype=wp.quat), cam_bodyid: wp.array(dtype=int), - cam_quat: wp.array(dtype=wp.quat), + cam_quat: wp.array2d(dtype=wp.quat), cam_fovy: wp.array(dtype=float), cam_resolution: wp.array(dtype=wp.vec2i), cam_sensorsize: wp.array(dtype=wp.vec2), From 459dd0f67007235691bd5dc607201eaee67c2724 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 14 May 2025 10:44:02 +0200 Subject: [PATCH 70/74] fix issues after merge --- mujoco_warp/_src/passive.py | 12 ++++++------ mujoco_warp/_src/ray.py | 10 ++++++---- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index d5588122..88cc6a4e 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -132,8 +132,8 @@ def _box_fluid( opt_density: float, opt_viscosity: float, body_rootid: wp.array(dtype=int), - body_mass: wp.array(dtype=float), - body_inertia: wp.array(dtype=wp.vec3), + body_mass: wp.array2d(dtype=float), + body_inertia: wp.array2d(dtype=wp.vec3), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), ximat_in: wp.array2d(dtype=wp.mat33), @@ -175,8 +175,8 @@ def _box_fluid( density = opt_density > 0.0 if viscosity or density: - inertia = body_inertia[bodyid] - mass = body_mass[bodyid] + inertia = body_inertia[worldid, bodyid] + mass = body_mass[worldid, bodyid] scl = 6.0 / mass box0 = wp.sqrt(wp.max(MJ_MINVAL, inertia[1] + inertia[2] - inertia[0]) * scl) box1 = wp.sqrt(wp.max(MJ_MINVAL, inertia[0] + inertia[2] - inertia[1]) * scl) @@ -249,7 +249,7 @@ def _qfrc_passive( # Model: jnt_actgravcomp: wp.array(dtype=int), dof_jntid: wp.array(dtype=int), - dof_damping: wp.array(dtype=float), + dof_damping: wp.array2d(dtype=float), # Data in: qvel_in: wp.array2d(dtype=float), qfrc_spring_in: wp.array2d(dtype=float), @@ -268,7 +268,7 @@ def _qfrc_passive( qfrc_passive = qfrc_spring_in[worldid, dofid] # damper - qfrc_damper = -dof_damping[dofid] * qvel_in[worldid, dofid] + qfrc_damper = -dof_damping[worldid, dofid] * qvel_in[worldid, dofid] qfrc_damper_out[worldid, dofid] = qfrc_damper qfrc_passive += qfrc_damper diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index eb713526..7915bb11 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -425,7 +425,7 @@ def _ray_geom_with_mesh( nmeshface: int, geom_type: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), - geom_size: wp.array(dtype=wp.vec3), + geom_size: wp.array2d(dtype=wp.vec3), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), mesh_vert: wp.array(dtype=wp.vec3), @@ -435,9 +435,10 @@ def _ray_geom_with_mesh( geom_id: int, pnt: wp.vec3, vec: wp.vec3, + worldid: int ) -> DistanceWithId: type = geom_type[geom_id] - size = geom_size[geom_id] + size = geom_size[worldid,geom_id] # TODO(team): static loop unrolling to remove unnecessary branching if type == int(GeomType.PLANE.value): @@ -498,7 +499,7 @@ def _ray_all_geom( geom_dataid: wp.array(dtype=int), geom_group: wp.array(dtype=int), geom_matid: wp.array(dtype=int), - geom_size: wp.array(dtype=wp.vec3), + geom_size: wp.array2d(dtype=wp.vec3), geom_rgba: wp.array(dtype=wp.vec4), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), @@ -581,6 +582,7 @@ def _ray_all_geom( geom_id, local_pnt, local_vec, + worldid, ) cur_dist = result.dist else: @@ -613,7 +615,7 @@ def _ray_all_geom_kernel( geom_dataid: wp.array(dtype=int), geom_group: wp.array(dtype=int), geom_matid: wp.array(dtype=int), - geom_size: wp.array(dtype=wp.vec3), + geom_size: wp.array2d(dtype=wp.vec3), geom_rgba: wp.array(dtype=wp.vec4), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), From 3f08a8cd8a98a624a090d9a004b65850a10c9c44 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 14 May 2025 10:46:06 +0200 Subject: [PATCH 71/74] formatting --- 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 c4a94840..2345d69c 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -645,11 +645,11 @@ class Model: jnt_actfrclimited: does joint have actuator force limits (njnt,) jnt_solref: constraint solver reference: limit (nworld, njnt, mjNREF) jnt_solimp: constraint solver impedance: limit (nworld, njnt, mjNIMP) - jnt_pos: local anchor position (nworld,njnt, 3) + jnt_pos: local anchor position (nworld, njnt, 3) jnt_axis: local joint axis (njnt, 3) jnt_stiffness: stiffness coefficient (nworld, njnt) jnt_range: joint limits (nworld, njnt, 2) - jnt_actfrcrange: range of total actuator force (nworldnjnt, 2) + jnt_actfrcrange: range of total actuator force (nworld, njnt, 2) jnt_margin: min distance for limit detection (nworld, njnt) jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr jnt_limited_ball_adr: limited/ball jntadr @@ -663,7 +663,7 @@ class Model: dof_invweight0: diag. inverse inertia in qpos0 (nworld, nv) dof_frictionloss: dof friction loss (nworld, nv) dof_solimp: constraint solver impedance: frictionloss (nworld, nv, NIMP) - dof_solref: constraint solver reference: frictionloss (nworld,nv, NREF) + dof_solref: constraint solver reference: frictionloss (nworld, nv, NREF) dof_tri_row: np.tril_indices (mjm.nv)[0] dof_tri_col: np.tril_indices (mjm.nv)[1] geom_type: geometric type (mjtGeom) (ngeom,) From a549afc38ebb9c1f9b0708f25d8ab10ddd18d129 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 14 May 2025 10:47:21 +0200 Subject: [PATCH 72/74] formatting --- mujoco_warp/_src/ray.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index 7915bb11..8e0f3f2d 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -435,10 +435,10 @@ def _ray_geom_with_mesh( geom_id: int, pnt: wp.vec3, vec: wp.vec3, - worldid: int + worldid: int, ) -> DistanceWithId: type = geom_type[geom_id] - size = geom_size[worldid,geom_id] + size = geom_size[worldid, geom_id] # TODO(team): static loop unrolling to remove unnecessary branching if type == int(GeomType.PLANE.value): From ab87b5c0cf8195a4a5348e1b61aac4b3cb75bab3 Mon Sep 17 00:00:00 2001 From: Alain Denzler Date: Wed, 14 May 2025 11:27:23 +0200 Subject: [PATCH 73/74] updates to follow MJX --- mujoco_warp/_src/collision_box.py | 2 +- mujoco_warp/_src/collision_convex.py | 2 +- mujoco_warp/_src/collision_primitive.py | 8 +-- mujoco_warp/_src/constraint.py | 16 +++--- mujoco_warp/_src/io.py | 20 +++---- mujoco_warp/_src/passive.py | 4 +- mujoco_warp/_src/ray.py | 18 +++---- mujoco_warp/_src/ray_test.py | 2 +- mujoco_warp/_src/sensor.py | 70 ++++++++++++------------- mujoco_warp/_src/smooth.py | 14 ++--- mujoco_warp/_src/types.py | 40 +++++++------- 11 files changed, 97 insertions(+), 99 deletions(-) diff --git a/mujoco_warp/_src/collision_box.py b/mujoco_warp/_src/collision_box.py index 5c4f2038..0ff58080 100644 --- a/mujoco_warp/_src/collision_box.py +++ b/mujoco_warp/_src/collision_box.py @@ -191,7 +191,7 @@ def _box_box( # Model: geom_type: wp.array(dtype=int), geom_condim: wp.array(dtype=int), - geom_priority: wp.array2d(dtype=int), + geom_priority: wp.array(dtype=int), geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), diff --git a/mujoco_warp/_src/collision_convex.py b/mujoco_warp/_src/collision_convex.py index 43e077e5..2f326759 100644 --- a/mujoco_warp/_src/collision_convex.py +++ b/mujoco_warp/_src/collision_convex.py @@ -708,7 +708,7 @@ def gjk_epa_sparse( geom_type: wp.array(dtype=int), geom_condim: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), - geom_priority: wp.array2d(dtype=int), + geom_priority: wp.array(dtype=int), geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), diff --git a/mujoco_warp/_src/collision_primitive.py b/mujoco_warp/_src/collision_primitive.py index 436ee4a3..ee60523d 100644 --- a/mujoco_warp/_src/collision_primitive.py +++ b/mujoco_warp/_src/collision_primitive.py @@ -1267,7 +1267,7 @@ def plane_cylinder( def contact_params( # Model: geom_condim: wp.array(dtype=int), - geom_priority: wp.array2d(dtype=int), + geom_priority: wp.array(dtype=int), geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), @@ -1303,8 +1303,8 @@ def contact_params( g1 = geoms[0] g2 = geoms[1] - p1 = geom_priority[worldid, g1] - p2 = geom_priority[worldid, g2] + p1 = geom_priority[g1] + p2 = geom_priority[g2] solmix1 = geom_solmix[worldid, g1] solmix2 = geom_solmix[worldid, g2] @@ -2380,7 +2380,7 @@ def _primitive_narrowphase( geom_type: wp.array(dtype=int), geom_condim: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), - geom_priority: wp.array2d(dtype=int), + geom_priority: wp.array(dtype=int), geom_solmix: wp.array2d(dtype=float), geom_solref: wp.array2d(dtype=wp.vec2), geom_solimp: wp.array2d(dtype=vec5), diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index 4341f505..b1ab3d37 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -107,7 +107,7 @@ def _efc_equality_connect( eq_objtype: wp.array(dtype=int), eq_solref: wp.array2d(dtype=wp.vec2), eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array(dtype=vec11), + eq_data: wp.array2d(dtype=vec11), eq_connect_adr: wp.array(dtype=int), # Data in: nefc_in: wp.array(dtype=int), @@ -142,7 +142,7 @@ def _efc_equality_connect( necid = wp.atomic_add(ne_connect_out, 0, 3) efcid = nefc_in[0] + necid - data = eq_data[i_eq] + data = eq_data[worldid, i_eq] anchor1 = wp.vec3f(data[0], data[1], data[2]) anchor2 = wp.vec3f(data[3], data[4], data[5]) @@ -239,7 +239,7 @@ def _efc_equality_joint( eq_obj2id: wp.array(dtype=int), eq_solref: wp.array2d(dtype=wp.vec2), eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array(dtype=vec11), + eq_data: wp.array2d(dtype=vec11), eq_jnt_adr: wp.array(dtype=int), # Data in: ne_connect_in: wp.array(dtype=int), @@ -272,7 +272,7 @@ def _efc_equality_joint( jntid_1 = eq_obj1id[i_eq] jntid_2 = eq_obj2id[i_eq] - data = eq_data[i_eq] + data = eq_data[worldid, i_eq] dofadr1 = jnt_dofadr[jntid_1] qposadr1 = jnt_qposadr[jntid_1] efc_J_out[efcid, dofadr1] = 1.0 @@ -330,7 +330,7 @@ def _efc_equality_tendon( eq_obj2id: wp.array(dtype=int), eq_solref: wp.array2d(dtype=wp.vec2), eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array(dtype=vec11), + eq_data: wp.array2d(dtype=vec11), eq_ten_adr: wp.array(dtype=int), tendon_length0: wp.array2d(dtype=float), tendon_invweight0: wp.array2d(dtype=float), @@ -368,7 +368,7 @@ def _efc_equality_tendon( obj1id = eq_obj1id[eqid] obj2id = eq_obj2id[eqid] - data = eq_data[eqid] + data = eq_data[worldid,eqid] solref = eq_solref[worldid, eqid] solimp = eq_solimp[worldid, eqid] pos1 = ten_length_in[worldid, obj1id] - tendon_length0[worldid, obj1id] @@ -498,7 +498,7 @@ def _efc_equality_weld( eq_objtype: wp.array(dtype=int), eq_solref: wp.array2d(dtype=wp.vec2), eq_solimp: wp.array2d(dtype=vec5), - eq_data: wp.array(dtype=vec11), + eq_data: wp.array2d(dtype=vec11), eq_wld_adr: wp.array(dtype=int), # Data in: ne_connect_in: wp.array(dtype=int), @@ -539,7 +539,7 @@ def _efc_equality_weld( obj1id = eq_obj1id[i_eq] obj2id = eq_obj2id[i_eq] - data = eq_data[i_eq] + data = eq_data[worldid, i_eq] anchor1 = wp.vec3(data[0], data[1], data[2]) anchor2 = wp.vec3(data[3], data[4], data[5]) relpose = wp.quat(data[6], data[7], data[8], data[9]) diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py index b0ddd999..8df48a2d 100644 --- a/mujoco_warp/_src/io.py +++ b/mujoco_warp/_src/io.py @@ -340,7 +340,7 @@ def create_nmodel_batched_array(mjm_array, dtype): body_invweight0=create_nmodel_batched_array(mjm.body_invweight0, dtype=float), body_contype=wp.array(mjm.body_contype, dtype=int), body_conaffinity=wp.array(mjm.body_conaffinity, dtype=int), - body_gravcomp=wp.array(mjm.body_gravcomp, dtype=float), + body_gravcomp=create_nmodel_batched_array(mjm.body_gravcomp, dtype=float), jnt_type=wp.array(mjm.jnt_type, dtype=int), jnt_qposadr=wp.array(mjm.jnt_qposadr, dtype=int), jnt_dofadr=wp.array(mjm.jnt_dofadr, dtype=int), @@ -350,7 +350,7 @@ def create_nmodel_batched_array(mjm_array, dtype): jnt_solref=create_nmodel_batched_array(mjm.jnt_solref, dtype=wp.vec2), jnt_solimp=create_nmodel_batched_array(mjm.jnt_solimp, dtype=types.vec5), jnt_pos=create_nmodel_batched_array(mjm.jnt_pos, dtype=wp.vec3), - jnt_axis=wp.array(mjm.jnt_axis, dtype=wp.vec3), + jnt_axis=create_nmodel_batched_array(mjm.jnt_axis, dtype=wp.vec3), jnt_stiffness=create_nmodel_batched_array(mjm.jnt_stiffness, dtype=float), jnt_range=create_nmodel_batched_array(mjm.jnt_range, dtype=float), jnt_actfrcrange=create_nmodel_batched_array(mjm.jnt_actfrcrange, dtype=wp.vec2), @@ -386,20 +386,20 @@ def create_nmodel_batched_array(mjm_array, dtype): geom_bodyid=wp.array(mjm.geom_bodyid, dtype=int), geom_dataid=wp.array(mjm.geom_dataid, dtype=int), geom_group=wp.array(mjm.geom_group, dtype=int), - geom_matid=wp.array(mjm.geom_matid, dtype=int), - geom_priority=create_nmodel_batched_array(mjm.geom_priority, dtype=int), + geom_matid=create_nmodel_batched_array(mjm.geom_matid, dtype=int), + geom_priority=wp.array(mjm.geom_priority, dtype=int), geom_solmix=create_nmodel_batched_array(mjm.geom_solmix, dtype=float), geom_solref=create_nmodel_batched_array(mjm.geom_solref, dtype=wp.vec2), geom_solimp=create_nmodel_batched_array(mjm.geom_solimp, dtype=types.vec5), geom_size=create_nmodel_batched_array(mjm.geom_size, dtype=wp.vec3), - geom_aabb=create_nmodel_batched_array(mjm.geom_aabb, dtype=wp.vec3), + geom_aabb=wp.array(mjm.geom_aabb, dtype=wp.vec3), geom_rbound=create_nmodel_batched_array(mjm.geom_rbound, dtype=float), geom_pos=create_nmodel_batched_array(mjm.geom_pos, dtype=wp.vec3), geom_quat=create_nmodel_batched_array(mjm.geom_quat, dtype=wp.quat), geom_friction=create_nmodel_batched_array(mjm.geom_friction, dtype=wp.vec3), geom_margin=create_nmodel_batched_array(mjm.geom_margin, dtype=float), geom_gap=create_nmodel_batched_array(mjm.geom_gap, dtype=float), - geom_rgba=wp.array(mjm.geom_rgba, dtype=wp.vec4), + geom_rgba=create_nmodel_batched_array(mjm.geom_rgba, dtype=wp.vec4), site_bodyid=wp.array(mjm.site_bodyid, dtype=int), site_pos=create_nmodel_batched_array(mjm.site_pos, dtype=wp.vec3), site_quat=create_nmodel_batched_array(mjm.site_quat, dtype=wp.quat), @@ -433,7 +433,7 @@ def create_nmodel_batched_array(mjm_array, dtype): eq_active0=wp.array(mjm.eq_active0, dtype=bool), eq_solref=create_nmodel_batched_array(mjm.eq_solref, dtype=wp.vec2), eq_solimp=create_nmodel_batched_array(mjm.eq_solimp, dtype=types.vec5), - eq_data=wp.array(mjm.eq_data, dtype=types.vec11), + eq_data=create_nmodel_batched_array(mjm.eq_data, dtype=types.vec11), # pre-compute indices of equality constraints eq_connect_adr=wp.array(np.nonzero(mjm.eq_type == types.EqType.CONNECT.value)[0], dtype=int), eq_wld_adr=wp.array(np.nonzero(mjm.eq_type == types.EqType.WELD.value)[0], dtype=int), @@ -487,7 +487,7 @@ def create_nmodel_batched_array(mjm_array, dtype): tendon_length0=create_nmodel_batched_array(mjm.tendon_length0, dtype=float), tendon_invweight0=create_nmodel_batched_array(mjm.tendon_invweight0, dtype=float), wrap_objid=wp.array(mjm.wrap_objid, dtype=int), - wrap_prm=create_nmodel_batched_array(mjm.wrap_prm, dtype=float), + wrap_prm=wp.array(mjm.wrap_prm, dtype=float), wrap_type=wp.array(mjm.wrap_type, dtype=int), tendon_jnt_adr=wp.array(tendon_jnt_adr, dtype=int), tendon_site_adr=wp.array(tendon_site_adr, dtype=int), @@ -505,7 +505,7 @@ def create_nmodel_batched_array(mjm_array, dtype): sensor_refid=wp.array(mjm.sensor_refid, dtype=int), sensor_dim=wp.array(mjm.sensor_dim, dtype=int), sensor_adr=wp.array(mjm.sensor_adr, dtype=int), - sensor_cutoff=create_nmodel_batched_array(mjm.sensor_cutoff, dtype=float), + sensor_cutoff=wp.array(mjm.sensor_cutoff, dtype=float), sensor_pos_adr=wp.array( np.nonzero(mjm.sensor_needstage == mujoco.mjtStage.mjSTAGE_POS)[0], dtype=int, @@ -532,7 +532,7 @@ def create_nmodel_batched_array(mjm_array, dtype): mujoco.mjtSensor.mjSENS_FRAMEANGACC, ], ).any(), - mat_rgba=wp.array(mjm.mat_rgba, dtype=wp.vec4), + mat_rgba=create_nmodel_batched_array(mjm.mat_rgba, dtype=wp.vec4), ) return m diff --git a/mujoco_warp/_src/passive.py b/mujoco_warp/_src/passive.py index 88cc6a4e..e5adacee 100644 --- a/mujoco_warp/_src/passive.py +++ b/mujoco_warp/_src/passive.py @@ -103,7 +103,7 @@ def _gravity_force( body_parentid: wp.array(dtype=int), body_rootid: wp.array(dtype=int), body_mass: wp.array2d(dtype=float), - body_gravcomp: wp.array(dtype=float), + body_gravcomp: wp.array2d(dtype=float), dof_bodyid: wp.array(dtype=int), # Data in: xipos_in: wp.array2d(dtype=wp.vec3), @@ -114,7 +114,7 @@ def _gravity_force( ): worldid, bodyid, dofid = wp.tid() bodyid += 1 # skip world body - gravcomp = body_gravcomp[bodyid] + gravcomp = body_gravcomp[worldid, bodyid] if gravcomp: force = -opt_gravity * body_mass[worldid, bodyid] * gravcomp diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index 8e0f3f2d..f1f44dee 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -498,15 +498,15 @@ def _ray_all_geom( geom_bodyid: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), geom_group: wp.array(dtype=int), - geom_matid: wp.array(dtype=int), + geom_matid: wp.array2d(dtype=int), geom_size: wp.array2d(dtype=wp.vec3), - geom_rgba: wp.array(dtype=wp.vec4), + geom_rgba: wp.array2d(dtype=wp.vec4), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), mesh_vert: wp.array(dtype=wp.vec3), mesh_faceadr: wp.array(dtype=int), mesh_face: wp.array(dtype=wp.vec3i), - mat_rgba: wp.array(dtype=wp.vec4), + 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), @@ -547,11 +547,11 @@ def _ray_all_geom( geom_filter = geom_filter and (geomgroup[group] != 0) # RGBA filter - matid = geom_matid[geom_id] - geom_alpha = geom_rgba[geom_id][3] + matid = geom_matid[worldid, geom_id] + geom_alpha = geom_rgba[worldid,geom_id][3] mat_alpha = wp.float32(0.0) if matid != -1: - mat_alpha = mat_rgba[matid][3] + mat_alpha = mat_rgba[worldid,matid][3] # Geom is visible if either: # 1. No material and non-zero geom alpha, or @@ -614,15 +614,15 @@ def _ray_all_geom_kernel( geom_bodyid: wp.array(dtype=int), geom_dataid: wp.array(dtype=int), geom_group: wp.array(dtype=int), - geom_matid: wp.array(dtype=int), + geom_matid: wp.array2d(dtype=int), geom_size: wp.array2d(dtype=wp.vec3), - geom_rgba: wp.array(dtype=wp.vec4), + geom_rgba: wp.array2d(dtype=wp.vec4), mesh_vertadr: wp.array(dtype=int), mesh_vertnum: wp.array(dtype=int), mesh_vert: wp.array(dtype=wp.vec3), mesh_faceadr: wp.array(dtype=int), mesh_face: wp.array(dtype=wp.vec3i), - mat_rgba: wp.array(dtype=wp.vec4), + 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), diff --git a/mujoco_warp/_src/ray_test.py b/mujoco_warp/_src/ray_test.py index 4bddd676..61655320 100644 --- a/mujoco_warp/_src/ray_test.py +++ b/mujoco_warp/_src/ray_test.py @@ -268,7 +268,7 @@ def test_ray_invisible(self): mjm, mjd, m, d = test_util.fixture("ray.xml") # nothing hit with transparent geoms - m.geom_rgba = wp.array([wp.vec4(0.0, 0.0, 0.0, 0.0)], dtype=wp.vec4) + m.geom_rgba = wp.array2d([[wp.vec4(0.0, 0.0, 0.0, 0.0)]], dtype=wp.vec4) mujoco.mj_forward(mjm, mjd) pnt = wp.array([wp.vec3(2.0, 1.0, 3.0)], dtype=wp.vec3) diff --git a/mujoco_warp/_src/sensor.py b/mujoco_warp/_src/sensor.py index 660a836a..2a5f43c0 100644 --- a/mujoco_warp/_src/sensor.py +++ b/mujoco_warp/_src/sensor.py @@ -34,16 +34,15 @@ def _write_scalar( # Model: sensor_datatype: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array2d(dtype=float), + sensor_cutoff: wp.array(dtype=float), # In: sensorid: int, sensor: Any, - worldid: int, # Out: out: wp.array(dtype=float), ): adr = sensor_adr[sensorid] - cutoff = sensor_cutoff[worldid, sensorid] + cutoff = sensor_cutoff[sensorid] if cutoff > 0.0: datatype = sensor_datatype[sensorid] @@ -60,17 +59,16 @@ def _write_vector( # Model: sensor_datatype: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array2d(dtype=float), + sensor_cutoff: wp.array(dtype=float), # In: sensorid: int, sensordim: int, sensor: Any, - worldid: int, # Out: out: wp.array(dtype=float), ): adr = sensor_adr[sensorid] - cutoff = sensor_cutoff[worldid, sensorid] + cutoff = sensor_cutoff[sensorid] if cutoff > 0.0: datatype = sensor_datatype[sensorid] @@ -374,7 +372,7 @@ def _sensor_pos( sensor_reftype: wp.array(dtype=int), sensor_refid: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array2d(dtype=float), + sensor_cutoff: wp.array(dtype=float), sensor_pos_adr: wp.array(dtype=int), # Data in: time_in: wp.array(dtype=float), @@ -407,19 +405,19 @@ def _sensor_pos( vec2 = _cam_projection( cam_fovy, cam_resolution, cam_sensorsize, cam_intrinsic, site_xpos_in, cam_xpos_in, cam_xmat_in, worldid, objid, refid ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 2, vec2, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 2, vec2, out) elif sensortype == int(SensorType.JOINTPOS.value): val = _joint_pos(jnt_qposadr, qpos_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) elif sensortype == int(SensorType.TENDONPOS.value): val = _tendon_pos(ten_length_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) elif sensortype == int(SensorType.ACTUATORPOS.value): val = _actuator_pos(actuator_length_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) elif sensortype == int(SensorType.BALLQUAT.value): quat = _ball_quat(jnt_qposadr, qpos_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, out) elif sensortype == int(SensorType.FRAMEPOS.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -441,7 +439,7 @@ def _sensor_pos( refid, reftype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif ( sensortype == int(SensorType.FRAMEXAXIS.value) or sensortype == int(SensorType.FRAMEYAXIS.value) @@ -459,7 +457,7 @@ def _sensor_pos( vec3 = _frame_axis( ximat_in, xmat_in, geom_xmat_in, site_xmat_in, cam_xmat_in, worldid, objid, objtype, refid, reftype, axis ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.FRAMEQUAT.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -479,13 +477,13 @@ def _sensor_pos( refid, reftype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 4, quat, out) elif sensortype == int(SensorType.SUBTREECOM.value): vec3 = _subtree_com(subtree_com_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.CLOCK.value): val = _clock(time_in, worldid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) @event_scope @@ -860,7 +858,7 @@ def _sensor_vel( sensor_reftype: wp.array(dtype=int), sensor_refid: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array2d(dtype=float), + sensor_cutoff: wp.array(dtype=float), sensor_vel_adr: wp.array(dtype=int), # Data in: qvel_in: wp.array2d(dtype=float), @@ -891,22 +889,22 @@ def _sensor_vel( if sensortype == int(SensorType.VELOCIMETER.value): vec3 = _velocimeter(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.GYRO.value): vec3 = _gyro(site_bodyid, site_xmat_in, cvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.JOINTVEL.value): val = _joint_vel(jnt_dofadr, qvel_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) elif sensortype == int(SensorType.TENDONVEL.value): val = _tendon_vel(ten_velocity_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) elif sensortype == int(SensorType.ACTUATORVEL.value): val = _actuator_vel(actuator_velocity_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) elif sensortype == int(SensorType.BALLANGVEL.value): vec3 = _ball_ang_vel(jnt_dofadr, qvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.FRAMELINVEL.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -934,7 +932,7 @@ def _sensor_vel( refid, reftype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_linvel, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_linvel, out) elif sensortype == int(SensorType.FRAMEANGVEL.value): objtype = sensor_objtype[sensorid] refid = sensor_refid[sensorid] @@ -962,13 +960,13 @@ def _sensor_vel( refid, reftype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_angvel, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, frame_angvel, out) elif sensortype == int(SensorType.SUBTREELINVEL.value): vec3 = _subtree_linvel(subtree_linvel_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.SUBTREEANGMOM.value): vec3 = _subtree_angmom(subtree_angmom_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) @event_scope @@ -1200,7 +1198,7 @@ def _sensor_acc( sensor_objtype: wp.array(dtype=int), sensor_objid: wp.array(dtype=int), sensor_adr: wp.array(dtype=int), - sensor_cutoff: wp.array2d(dtype=float), + sensor_cutoff: wp.array(dtype=float), sensor_acc_adr: wp.array(dtype=int), # Data in: xpos_in: wp.array2d(dtype=wp.vec3), @@ -1228,19 +1226,19 @@ def _sensor_acc( vec3 = _accelerometer( body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cvel_in, cacc_in, worldid, objid ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.FORCE.value): vec3 = _force(site_bodyid, site_xmat_in, cfrc_int_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.TORQUE.value): vec3 = _torque(body_rootid, site_bodyid, site_xpos_in, site_xmat_in, subtree_com_in, cfrc_int_in, worldid, objid) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.ACTUATORFRC.value): val = _actuator_force(actuator_force_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) elif sensortype == int(SensorType.JOINTACTFRC.value): val = _joint_actuator_force(jnt_dofadr, qfrc_actuator_in, worldid, objid) - _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, worldid, out) + _write_scalar(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, val, out) elif sensortype == int(SensorType.FRAMELINACC.value): objtype = sensor_objtype[sensorid] vec3 = _framelinacc( @@ -1260,7 +1258,7 @@ def _sensor_acc( objid, objtype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) elif sensortype == int(SensorType.FRAMEANGACC.value): objtype = sensor_objtype[sensorid] vec3 = _frameangacc( @@ -1272,7 +1270,7 @@ def _sensor_acc( objid, objtype, ) - _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, worldid, out) + _write_vector(sensor_datatype, sensor_adr, sensor_cutoff, sensorid, 3, vec3, out) @event_scope diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 87b2e96d..593b1ad8 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -71,7 +71,7 @@ def _kinematics_level( jnt_type: wp.array(dtype=int), jnt_qposadr: wp.array(dtype=int), jnt_pos: wp.array2d(dtype=wp.vec3), - jnt_axis: wp.array(dtype=wp.vec3), + jnt_axis: wp.array2d(dtype=wp.vec3), # Data in: qpos_in: wp.array2d(dtype=float), xpos_in: wp.array2d(dtype=wp.vec3), @@ -105,7 +105,7 @@ def _kinematics_level( xpos = wp.vec3(qpos[qadr], qpos[qadr + 1], qpos[qadr + 2]) xquat = wp.quat(qpos[qadr + 3], qpos[qadr + 4], qpos[qadr + 5], qpos[qadr + 6]) xanchor_out[worldid, jntadr] = xpos - xaxis_out[worldid, jntadr] = jnt_axis[jntadr] + xaxis_out[worldid, jntadr] = jnt_axis[worldid, jntadr] else: # regular or no joints # apply fixed translation and rotation relative to parent @@ -116,7 +116,7 @@ def _kinematics_level( for _ in range(jntnum): qadr = jnt_qposadr[jntadr] jnt_type_ = jnt_type[jntadr] - jnt_axis_ = jnt_axis[jntadr] + jnt_axis_ = jnt_axis[worldid, jntadr] xanchor = math.rot_vec_quat(jnt_pos[worldid, jntadr], xquat) + xpos xaxis = math.rot_vec_quat(jnt_axis_, xquat) @@ -1144,7 +1144,7 @@ def _cfrc_ext_equality( eq_obj1id: wp.array(dtype=int), eq_obj2id: wp.array(dtype=int), eq_objtype: wp.array(dtype=int), - eq_data: wp.array(dtype=vec11), + eq_data: wp.array2d(dtype=vec11), # Data in: ne_connect_in: wp.array(dtype=int), ne_weld_in: wp.array(dtype=int), @@ -1182,7 +1182,7 @@ def _cfrc_ext_equality( worldid = efc_worldid_in[efcid] id = efc_id_in[efcid] - eq_data_ = eq_data[id] + eq_data_ = eq_data[worldid,id] body_semantic = eq_objtype[id] == wp.static(ObjType.BODY.value) obj1 = eq_obj1id[id] @@ -1905,7 +1905,7 @@ def _joint_tendon( jnt_qposadr: wp.array(dtype=int), jnt_dofadr: wp.array(dtype=int), wrap_objid: wp.array(dtype=int), - wrap_prm: wp.array2d(dtype=float), + wrap_prm: wp.array(dtype=float), tendon_jnt_adr: wp.array(dtype=int), wrap_jnt_adr: wp.array(dtype=int), # Data in: @@ -1920,7 +1920,7 @@ def _joint_tendon( wrap_jnt_adr_ = wrap_jnt_adr[wrapid] wrap_objid_ = wrap_objid[wrap_jnt_adr_] - prm = wrap_prm[worldid, wrap_jnt_adr_] + prm = wrap_prm[wrap_jnt_adr_] # add to length L = prm * qpos_in[worldid, jnt_qposadr[wrap_objid_]] diff --git a/mujoco_warp/_src/types.py b/mujoco_warp/_src/types.py index 2345d69c..3d3dd3fc 100644 --- a/mujoco_warp/_src/types.py +++ b/mujoco_warp/_src/types.py @@ -636,7 +636,7 @@ class Model: body_invweight0: mean inv inert in qpos0 (trn, rot) (nworld, nbody, 2) body_contype: OR over all geom contypes (nbody,) body_conaffinity: OR over all geom conaffinities (nbody,) - body_gravcomp: antigravity force, units of body weight (nbody,) + body_gravcomp: antigravity force, units of body weight (nworld, nbody) jnt_type: type of joint (mjtJoint) (njnt,) jnt_qposadr: start addr in 'qpos' for joint's data (njnt,) jnt_dofadr: start addr in 'qvel' for joint's data (njnt,) @@ -646,7 +646,7 @@ class Model: jnt_solref: constraint solver reference: limit (nworld, njnt, mjNREF) jnt_solimp: constraint solver impedance: limit (nworld, njnt, mjNIMP) jnt_pos: local anchor position (nworld, njnt, 3) - jnt_axis: local joint axis (njnt, 3) + jnt_axis: local joint axis (nworld, njnt, 3) jnt_stiffness: stiffness coefficient (nworld, njnt) jnt_range: joint limits (nworld, njnt, 2) jnt_actfrcrange: range of total actuator force (nworld, njnt, 2) @@ -673,20 +673,20 @@ class Model: geom_bodyid: id of geom's body (ngeom,) geom_dataid: id of geom's mesh/hfield; -1: none (ngeom,) geom_group: geom group inclusion/exclusion mask (ngeom,) - geom_matid: material id for rendering (ngeom,) - geom_priority: geom contact priority (nworld, ngeom,) + geom_matid: material id for rendering (nworld, ngeom,) + geom_priority: geom contact priority (ngeom,) geom_solmix: mixing coef for solref/imp in geom pair (nworld, ngeom,) geom_solref: constraint solver reference: contact (nworld, ngeom, mjNREF) geom_solimp: constraint solver impedance: contact (nworld, ngeom, mjNIMP) geom_size: geom-specific size parameters (ngeom, 3) - geom_aabb: bounding box, (center, size) (nworld, ngeom, 6) + geom_aabb: bounding box, (center, size) (ngeom, 6) geom_rbound: radius of bounding sphere (nworld, ngeom,) geom_pos: local position offset rel. to body (nworld, ngeom, 3) geom_quat: local orientation offset rel. to body (nworld, ngeom, 4) geom_friction: friction for (slide, spin, roll) (nworld, ngeom, 3) geom_margin: detect contact if dist Date: Wed, 14 May 2025 11:27:45 +0200 Subject: [PATCH 74/74] formatting --- mujoco_warp/_src/constraint.py | 2 +- mujoco_warp/_src/ray.py | 4 ++-- mujoco_warp/_src/smooth.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/mujoco_warp/_src/constraint.py b/mujoco_warp/_src/constraint.py index b1ab3d37..6a09bdfe 100644 --- a/mujoco_warp/_src/constraint.py +++ b/mujoco_warp/_src/constraint.py @@ -368,7 +368,7 @@ def _efc_equality_tendon( obj1id = eq_obj1id[eqid] obj2id = eq_obj2id[eqid] - data = eq_data[worldid,eqid] + data = eq_data[worldid, eqid] solref = eq_solref[worldid, eqid] solimp = eq_solimp[worldid, eqid] pos1 = ten_length_in[worldid, obj1id] - tendon_length0[worldid, obj1id] diff --git a/mujoco_warp/_src/ray.py b/mujoco_warp/_src/ray.py index f1f44dee..f877d1fd 100644 --- a/mujoco_warp/_src/ray.py +++ b/mujoco_warp/_src/ray.py @@ -548,10 +548,10 @@ def _ray_all_geom( # RGBA filter matid = geom_matid[worldid, geom_id] - geom_alpha = geom_rgba[worldid,geom_id][3] + geom_alpha = geom_rgba[worldid, geom_id][3] mat_alpha = wp.float32(0.0) if matid != -1: - mat_alpha = mat_rgba[worldid,matid][3] + mat_alpha = mat_rgba[worldid, matid][3] # Geom is visible if either: # 1. No material and non-zero geom alpha, or diff --git a/mujoco_warp/_src/smooth.py b/mujoco_warp/_src/smooth.py index 593b1ad8..b7db29d9 100644 --- a/mujoco_warp/_src/smooth.py +++ b/mujoco_warp/_src/smooth.py @@ -1182,7 +1182,7 @@ def _cfrc_ext_equality( worldid = efc_worldid_in[efcid] id = efc_id_in[efcid] - eq_data_ = eq_data[worldid,id] + eq_data_ = eq_data[worldid, id] body_semantic = eq_objtype[id] == wp.static(ObjType.BODY.value) obj1 = eq_obj1id[id]