problem with equality_constraint

[kassasin]

I got an error when create a problem with constrained robot:

Traceback (most recent call last):
  File "/workspace/mjinx_example/examples/grmini_t1.py", line 127, in <module>
    opt_solution, solver_data = solve_jit(q, solver_data, problem_data)
  File "/workspace/mjinx/solvers/_base.py", line 106, in solve
    return self.solve_from_data(solver_data, problem_data, model_data)
  File "/workspace/mjinx/solvers/_local_ik.py", line 350, in solve_from_data
    P, c, A, b, G, h = self.__compute_qp_matrices(problem_data, model_data)
  File "/workspace/mjinx/solvers/_local_ik.py", line 323, in __compute_qp_matrices
    H, c = process_soft_constraint(component)
  File "/workspace/mjinx/solvers/_local_ik.py", line 271, in process_soft_constraint
    jacobian = constraint.compute_jacobian(model_data)
  File "/workspace/mjinx/components/constraints/_equality_constraint.py", line 50, in compute_jacobian
    return data.efc_J[data.efc_type == mj.mjtConstraint.mjCNSTR_EQUALITY, :][self.mask_idxs, :]
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

<equality>
    <connect name="right_foot_motor1_connect"
             site1="right_foot_motor1_connect"
             site2="right_foot_motor1_bar_connect"
             solimp="0.9 0.95 0.001"
             solref="0.01 1"/>
    <connect name="right_foot_motor2_connect"
             site1="right_foot_motor2_connect"
             site2="right_foot_motor2_bar_connect"
             solimp="0.9 0.95 0.001"
             solref="0.01 1"/>
    <connect name="left_foot_motor1_connect"
             site1="left_foot_motor1_connect"
             site2="left_foot_motor1_bar_connect"
             solimp="0.9 0.95 0.001"
             solref="0.01 1"/>
    <connect name="left_foot_motor2_connect"
             site1="left_foot_motor2_connect"
             site2="left_foot_motor2_bar_connect"
             solimp="0.9 0.95 0.001"
             solref="0.01 1"/>
</equality>

[domrachev03]

Hey there, thank you for using my library!

The constraints specification looks perfectly fine. Could you share the example file grmini_t1.py, preferably with the model that you are launching? If this is not possible, please at least show how you initialise equality constraints and compile the problem.

[kassasin]

Hey there, thank you for using my library!

The constraints specification looks perfectly fine. Could you share the example file grmini_t1.py, preferably with the model that you are launching? If this is not possible, please at least show how you initialise equality constraints and compile the problem.
Thanks ! I will upload my script and robot model later!

[kassasin]

This is my robot xml file :

<mujoco model="grmini_t1">
  <compiler angle="radian" meshdir="meshes/"/>

  <!-- timestep 表示仿真的基础时间步长是 0.002 秒 -->
  <!-- integrator="Euler",欧拉积分是最快、最简单的积分方法，但数值稳定性最差,可选项还有 RK4（Runge-Kutta 四阶）和 implicit（隐式积分，稳定性更强） -->
  <!-- iterations="3",是用于求解接触和约束的 PGS（投影梯度法）迭代次数,10~50，但你用了 3 ⇒ 求解更快但精度较低 -->
  <!-- ls_iterations="6",表示 线搜索迭代次数，用于处理接触时的约束解稳定性 -->
  <option iterations="3" ls_iterations="6" timestep=".002" integrator="Euler">
    <!-- 禁用 Euler damping,在 Euler 模式下添加轻微阻尼以防震荡 -->
    <flag eulerdamp="disable"/>
  </option>

  <default>
    <default class="grmini_t1">
      <!-- 几何体（geom）的接触维度为 1（通常是滑动摩擦）。

      contype=0 和 conaffinity=0 表示该物体不参与碰撞（默认非碰撞几何体）。 -->
      <geom condim="1" contype="0" conaffinity="0"/>
      <!-- 所有关节默认具有摩擦损失为 1，阻尼为 1。 -->
      <!-- 所有关节默认使用比例控制器，kp=100，速度反馈增益 kv=5。

      inheritrange=1 表示从父类继承关节范围（range）。 -->
      <joint frictionloss="1" damping="1"/>
      <!-- 只对设置了 <position> 控制的关节生效 -->
      <position inheritrange="1" kp="100" kv ="5"/>

      <default class="hip">
        <default class="hip_pitch">
          <position kp="180"  kv ="10"/>
          <joint axis="0 1 0" range="-2.617 2.617" actuatorfrcrange="-95 95" armature="0.01017752004"/>
        </default>
        <default class="left_hip_roll">
          <position kp="120"  kv ="10"/>
          <joint axis="1 0 0" range="-0.261 1.57" actuatorfrcrange="-54 54" armature="0.025101925"/>
        </default>
        <!-- reverse -->
        <default class="right_hip_roll">
          <position kp="120"  kv ="10"/>
          <joint axis="1 0 0" range="-1.57 0.261" actuatorfrcrange="-54 54" armature="0.025101925"/>
        </default>
        <default class="hip_yaw">
          <position kp="90"  kv ="8"/>
          <joint axis="0 0 1" range="-2.617 2.617" actuatorfrcrange="-54 54" armature="0.01017752004"/>
        </default>
      </default>
      <default class="knee">
        <position kp="120"  kv ="8"/>
        <joint axis="0 1 0" range="-0.0872 2.356" actuatorfrcrange="-95 95" armature="0.025101925"/>
      </default>
      <default class="foot_motor">
        <default class="foot_motor2">
          <position kp="45"  kv ="2.5"/>
          <joint axis="-1 0 0" range="-0.436 0.436" actuatorfrcrange="-30 30" damping="0.5" armature="0.00721945"/>
        </default>
        <default class="foot_motor1">
          <position kp="45"  kv ="2.5"/>
          <joint axis="0 -1 0" range="-0.436 0.436" actuatorfrcrange="-30 30" damping="0.5" armature="0.00721945"/>
        </default>
        <default class="left_foot_motor1">
          <position kp="45"  kv ="2.5"/>
          <joint axis="0 1 0" range="-0.436 0.436" actuatorfrcrange="-30 30" damping="0.5" armature="0.00721945"/>
        </default>
      </default>
      <default class="waist_yaw">
        <position kp="90"  kv ="8"/>
        <joint axis="0 0 1" range="-2.617 2.617" actuatorfrcrange="-54 54" armature="0.01017752004"/>
      </default>
      <default class="shoulder">
        <default class="shoulder_pitch">
          <position kp="90"  kv ="8"/>
          <joint axis="0 1 0" range="-2.966 2.966" actuatorfrcrange="-54 54" armature="0.003609725"/>
        </default>
        <default class="left_shoulder_roll">
          <position kp="45"  kv ="2.5"/>
          <joint axis="1 0 0" range="-0.174 2.792" actuatorfrcrange="-30 30" armature="0.003609725"/>
        </default>
        <!-- reverse -->
        <default class="right_shoulder_roll">
          <position kp="45"  kv ="2.5"/>
          <joint axis="1 0 0" range="-2.792 0.174" actuatorfrcrange="-30 30" armature="0.003609725"/>
        </default>
        <default class="shoulder_yaw">
          <position kp="45"  kv ="2.5"/>
          <joint axis="0 0 1" range="-1.832 1.832" actuatorfrcrange="-30 30" armature="0.003609725"/>
        </default>
      </default>
      <default class="elbow">
        <position kp="45"  kv ="2.5"/>
        <joint axis="0 1 0" range="-0.349 1.658" actuatorfrcrange="-30 30" armature="0.003609725"/>
      </default>
      <default class="wrist">
        <position kp="45"  kv ="2.5"/>
        <joint damping=".2"/>
        <default class="wrist_yaw">
          <joint axis="0 0 -1" range="-1.832 1.832" actuatorfrcrange="-30 30" armature="0.003609725"/>
        </default>
      </default>

      <!-- group=2，属于视觉组（渲染用）。

      类型是 mesh 模型，密度为 0，表示不参与物理模拟。 -->
      <default class="visual">
        <geom group="2" type="mesh" density="0"/>
      </default>
      <default class="collision">
      <!-- group=3，用于碰撞检测。有颜色但不参与质量计算（density=0） -->
        <geom group="3" rgba="0.8 0.6 0 1" density="0"/>
        <default class="foot">
          <geom size="0.085 0.03 0.005"/>
        </default>
      </default>
      <!-- 用于添加红色 site 点（可用于脚底、目标、传感器等定位用途）。 -->
      <site group="5" rgba="1 0 0 1"/>
    </default>
  </default>

  <asset>
    <material name="silver" rgba="0.7 0.7 0.7 1"/>
    <material name="black" rgba="0.2 0.2 0.2 1"/>
    <mesh name="base_link" file="base_link.STL"/>
    <mesh name="waist_yaw_link" file="waist_yaw_link.STL"/>
    <mesh name="torso_link" file="torso_link.STL"/>
    <mesh name="imu_link" file="imu_link.STL"/>
    <mesh name="camera_link" file="camera_link.STL"/>

    <mesh name="left_upper_arm_pitch_link" file="left_upper_arm_pitch_link.STL"/>
    <mesh name="left_upper_arm_roll_link" file="left_upper_arm_roll_link.STL"/>
    <mesh name="left_upper_arm_yaw_link" file="left_upper_arm_yaw_link.STL"/>
    <mesh name="left_lower_arm_pitch_link" file="left_lower_arm_pitch_link.STL"/>
    <mesh name="left_hand_yaw_link" file="left_hand_yaw_link.STL"/>
    <mesh name="left_end_effector_link" file="left_end_effector_link.STL"/>

    <mesh name="right_upper_arm_pitch_link" file="right_upper_arm_pitch_link.STL"/>
    <mesh name="right_upper_arm_roll_link" file="right_upper_arm_roll_link.STL"/>
    <mesh name="right_upper_arm_yaw_link" file="right_upper_arm_yaw_link.STL"/>
    <mesh name="right_lower_arm_pitch_link" file="right_lower_arm_pitch_link.STL"/>
    <mesh name="right_hand_yaw_link" file="right_hand_yaw_link.STL"/>
    <mesh name="right_end_effector_link" file="right_end_effector_link.STL"/>

    <mesh name="left_thigh_pitch_link" file="left_thigh_pitch_link.STL"/>
    <mesh name="left_thigh_roll_link" file="left_thigh_roll_link.STL"/>
    <mesh name="left_thigh_yaw_link" file="left_thigh_yaw_link.STL"/>
    <mesh name="left_shank_pitch_base_link" file="left_shank_pitch_base_link.stl"/>
    <mesh name="left_foot_roll_link" file="left_foot_roll_link.STL"/>
    <mesh name="left_foot_pitch_link" file="left_foot_pitch_link.STL"/>
    <mesh name="left_shank_pitch_motor1_link" file="left_shank_pitch_motor1_link.stl"/>
    <mesh name="left_shank_pitch_motor2_link" file="left_shank_pitch_motor2_link.stl"/>
    <mesh name="left_shank_pitch_motor1_bar_link" file="left_shank_pitch_motor1_bar_link.stl"/>
    <mesh name="left_shank_pitch_motor2_bar_link" file="left_shank_pitch_motor2_bar_link.stl"/>

    <mesh name="right_thigh_pitch_link" file="right_thigh_pitch_link.STL"/>
    <mesh name="right_thigh_roll_link" file="right_thigh_roll_link.STL"/>
    <mesh name="right_thigh_yaw_link" file="right_thigh_yaw_link.STL"/>
    <mesh name="right_shank_pitch_base_link" file="right_shank_pitch_base_link.stl"/>
    <mesh name="right_foot_roll_link" file="right_foot_roll_link.STL"/>
    <mesh name="right_foot_pitch_link" file="right_foot_pitch_link.STL"/>
    <mesh name="right_shank_pitch_motor1_link" file="right_shank_pitch_motor1_link.stl"/>
    <mesh name="right_shank_pitch_motor2_link" file="right_shank_pitch_motor2_link.stl"/>
    <mesh name="right_shank_pitch_motor1_bar_link" file="right_shank_pitch_motor1_bar_link.stl"/>
    <mesh name="right_shank_pitch_motor2_bar_link" file="right_shank_pitch_motor2_bar_link.stl"/>
  </asset>

  <worldbody>
    <body name="base_link" pos="0 0 0.710" childclass="grmini_t1">
        <camera name="track" pos="1.734 -1.135 .35" xyaxes="0.552 0.834 -0.000 -0.170 0.112 0.979" mode="trackcom"/>
        <inertial pos="-0.00036119 0 -0.05865679" quat="1 0 -0.000399148 0" mass="3.136" diaginertia="0.01354598 0.01005144 0.00953678"/>
        <freejoint name="floating_base_joint"/>
        <geom class="visual" material="black" mesh="base_link"/>
        <site name="imu" size="0.01" pos="0.07751 0 -0.0536"/>
        <body name="waist_yaw_link" pos="0 0 0.015">
            <inertial pos="0.00392916 -0.00032484 0.196302" quat="0.999809 -0.00128852 0.0194931 0.000508804" mass="8.8555" diaginertia="0.122593 0.113694 0.0393338"/>
            <joint name="waist_yaw_joint" class="waist_yaw"/>
            <geom pos="0 0 0.2604" quat="1 0 0 0" type="mesh" contype="0" conaffinity="0" group="2" density="0" rgba="1 1 1 1" mesh="torso_link"/>
            <geom pos="0.07875 0 0.0536" quat="1 0 0 0" type="mesh" contype="0" conaffinity="0" group="2" density="0" rgba="1 1 1 1" mesh="imu_link"/>
            <geom pos="0.071504 0 0.4402" quat="0.965926 0 0.25882 0" type="mesh" contype="0" conaffinity="0" group="2" density="0" rgba="1 1 1 1" mesh="camera_link"/>
            <geom class="visual" mesh="waist_yaw_link"/>
            <geom pos="0 0 0.2" name="waist_yaw_link_collision" class="collision" type="box" size="0.07 0.1 0.28"/>
            <body name="left_upper_arm_pitch_link" pos="0 0.09365 0.2604">
                <inertial pos="-0.00010549 0.0460864 8.03e-06" quat="0.495115 0.500204 -0.50005 0.504586" mass="0.8842" diaginertia="0.00184165 0.00183028 0.000625398"/>
                <joint name="left_shoulder_pitch_joint" class="shoulder_pitch"/>
                <geom class="visual" mesh="left_upper_arm_pitch_link"/>
                <body name="left_upper_arm_roll_link" pos="0 0.082 0">
                    <inertial pos="0.00487758 0.00044474 -0.0431393" quat="0.708386 0.0373823 0.0395157 0.703726" mass="0.714" diaginertia="0.00181005 0.00172856 0.000400062"/>
                    <joint name="left_shoulder_roll_joint" class="left_shoulder_roll"/>
                    <geom class="visual" mesh="left_upper_arm_roll_link"/>
                    <geom pos="0 0 -0.2" name="left_arm_collision" class="collision" type="sphere" size="0.060"/>
                    <body name="left_upper_arm_yaw_link" pos="0 0 -0.105">
                        <inertial pos="-9.046e-05 -0.00031442 -0.0556749" quat="0.988512 -0.00409726 -0.000239432 0.151085" mass="0.914" diaginertia="0.00309435 0.00308521 0.000737412"/>
                        <joint name="left_shoulder_yaw_joint" class="shoulder_yaw"/>
                        <geom class="visual" mesh="left_upper_arm_yaw_link"/>
                        <body name="left_lower_arm_pitch_link" pos="0 0 -0.105" quat="0.707105 0 -0.707108 0">
                            <inertial pos="-0.00041353 0.00475782 -0.0429711" quat="0.998474 -0.0548672 -0.0019523 -0.00593138" mass="0.712" diaginertia="0.00180547 0.00172538 0.000400887"/>
                            <joint name="left_elbow_pitch_joint" class="elbow"/>
                            <geom class="visual" mesh="left_lower_arm_pitch_link"/>
                            <body name="left_hand_yaw_link" pos="0 0 -0.105">
                                <inertial pos="-9.51713e-05 1.29434e-06 -0.0246028" quat="0.451748 -0.0015129 -0.000673209 0.892144" mass="0.57701" diaginertia="0.00130576 0.00130295 0.000411831"/>
                                <joint name="left_wrist_yaw_joint" class="wrist_yaw"/>
                                <geom class="visual" mesh="left_hand_yaw_link"/>
                                <geom pos="0 0 -0.08" quat="1 0 0 0" type="mesh" contype="0" conaffinity="0" group="2" density="0" rgba="1 1 1 1" mesh="left_end_effector_link"/>
                                <site name="left_palm" pos="0 0 -0.08" size="0.01"/>
                                <geom pos="0 0 -0.1" name="left_hand_collision" class="collision" type="sphere" size="0.030"/>
                            </body>
                        </body>
                    </body>
                </body>
            </body>
            <body name="right_upper_arm_pitch_link" pos="0 -0.09365 0.2604">
                <inertial pos="-0.000108 -0.0460864 8.05e-06" quat="0.494174 0.488789 -0.51052 0.506209" mass="0.8842" diaginertia="0.00184155 0.00183027 0.000625685"/>
                <joint name="right_shoulder_pitch_joint" class="shoulder_pitch"/>
                <geom class="visual" mesh="right_upper_arm_pitch_link"/>
                <body name="right_upper_arm_roll_link" pos="0 -0.082 0">
                    <inertial pos="0.00511059 -0.00047879 -0.0434469" quat="0.699891 0.0420404 0.0396474 0.711908" mass="0.711" diaginertia="0.00179303 0.00172216 0.000390187"/>
                    <joint name="right_shoulder_roll_joint" class="right_shoulder_roll"/>
                    <geom class="visual" mesh="right_upper_arm_roll_link"/>
                    <geom pos="0 0 -0.2" name="right_arm_collision" class="collision" type="sphere" size="0.060"/>
                    <body name="right_upper_arm_yaw_link" pos="0 0 -0.105">
                        <inertial pos="-9.057e-05 0.00031679 -0.0556758" quat="0.989909 0.00414525 -7.3143e-05 -0.141647" mass="0.914" diaginertia="0.00309365 0.00308561 0.000737058"/>
                        <joint name="right_shoulder_yaw_joint" class="shoulder_yaw"/>
                        <geom class="visual" mesh="right_upper_arm_yaw_link"/>
                        <body name="right_lower_arm_pitch_link" pos="0 0 -0.105" quat="0.707105 0 -0.707108 0">
                            <inertial pos="-0.00034586 -0.00474975 -0.0431673" quat="0.998479 0.0549162 -0.00170368 0.00452907" mass="0.712" diaginertia="0.00180972 0.00172691 0.0003979"/>
                            <joint name="right_elbow_pitch_joint" class="elbow"/>
                            <geom class="visual" mesh="right_lower_arm_pitch_link"/>
                            <body name="right_hand_yaw_link" pos="0 0 -0.105">
                                <inertial pos="-9.52307e-05 0 -0.0245736" quat="0.445169 -0.00154029 -0.000693026 0.895445" mass="0.5767" diaginertia="0.00128997 0.00128717 0.000402601"/>
                                <joint name="right_wrist_yaw_joint" class="wrist_yaw"/>
                                <geom class="visual" mesh="right_hand_yaw_link"/>
                                <geom pos="0 0 -0.08" quat="1 0 0 0" type="mesh" contype="0" conaffinity="0" group="2" density="0" rgba="1 1 1 1" mesh="right_end_effector_link"/>
                                <site name="right_palm" pos="0 0 -0.08" size="0.01"/>
                                <geom pos="0 0 -0.1" name="right_hand_collision" class="collision" type="sphere" size="0.030"/>
                            </body>
                        </body>
                    </body>
                </body>
            </body>
        </body>
    <body name="left_thigh_pitch_link" pos="0 0.048521 -0.076847" quat="0.991445 -0.130526 0 0">
        <inertial pos="0.00469728 0.0370354 0.00024962" quat="0.586057 0.394562 -0.588457 0.393163" mass="1.3008" diaginertia="0.00238065 0.00207849 0.00131154"/>
        <joint name="left_hip_pitch_joint" class="hip_pitch"/>
        <geom class="visual" mesh="left_thigh_pitch_link"/>
        <body name="left_thigh_roll_link" pos="0 0.074 0" quat="0.991445 0.130526 0 0">
            <inertial pos="0.00183717 -4.3e-06 -0.0321936" quat="0.999749 3.9204e-05 0.0224234 -2.26631e-05" mass="1.359" diaginertia="0.00302075 0.00288693 0.00125139"/>
            <joint name="left_hip_roll_joint" class="left_hip_roll"/>
            <geom class="visual" mesh="left_thigh_roll_link"/>
            <geom name="left_thigh" class="collision" type="capsule" size="0.07" fromto="0.0 -0.01 0 0.0 -0.01 -0.2"/>
            <body name="left_thigh_yaw_link" pos="0 0 -0.0875">
                <inertial pos="0.00216478 -0.00975305 -0.0835706" quat="0.591547 -0.0398546 0.00911377 0.805234" mass="3.452" diaginertia="0.016329 0.015951 0.00542247"/>
                <joint name="left_hip_yaw_joint" class="hip_yaw"/>
                <geom class="visual" mesh="left_thigh_yaw_link"/>
                <body name="left_shank_pitch_base_link" pos="0 0 -0.2031">
                    <inertial pos="0.00278302 0.00416562 -0.135619" quat="0.923116 -0.00745391 0.0122662 0.384253" mass="2.335" diaginertia="0.0138036 0.0137193 0.00176335"/>
                    <joint name="left_knee_pitch_joint" class="knee"/>
                    <geom class="visual" mesh="left_shank_pitch_base_link"/>
                    <geom name="left_shin" class="collision" type="capsule" size="0.07" fromto="0.02 0 0 0.02 0 -0.25"/>
                    <body name="left_foot_roll_link" pos="0 0 -0.28">
                        <inertial pos="0 0.001805 0" quat="0.707107 0 0 0.707107" mass="0.11173" diaginertia="2.876e-05 2.445e-05 1.162e-05"/>
                        <joint name="left_ankle_roll_joint" axis="1 0 0" damping="0.5" frictionloss="1"/>
                        <geom class="visual" mesh="left_foot_roll_link"/>
                        <body name="left_foot_pitch_link">
                            <site name="left_foot" rgba="1 0 0 1"/>
                            <site name="left_foot_motor1_connect" pos="0.05 0 0"/>
                            <site name="left_foot_motor2_connect" pos="0 0.029 0"/>
                            <inertial pos="0.023434 -0.00015234 -0.030677" quat="-0.000812256 0.705674 -0.00501102 0.708518" mass="0.391" diaginertia="0.000895276 0.00083368 0.000140403"/>
                            <joint name="left_ankle_pitch_joint" axis="0 1 0" damping="0.5" frictionloss="1"/>
                            <geom class="visual" mesh="left_foot_pitch_link"/>
                            <geom name="left_foot" class="foot" pos="0.035 0 -0.035" size="0.10 0.04 0.008" type="box"/>
                        </body>
                    </body>
                    <body name="left_shank_pitch_motor2_link" pos="0.0205 0.0 -0.195">
                        <inertial pos="0 0 0" quat="0.5 0.5 -0.5 0.5" mass="0.05" diaginertia="1e-5 1e-5 1e-5"/>
                        <joint name="left_ankle_motor2_joint" class="foot_motor2"/>
                        <geom class="visual" mesh="left_shank_pitch_motor2_link"/>
                        <body name="left_shank_pitch_motor2_bar_link" pos="0.0245 0.028 0">
                            <site name="left_foot_motor2_bar_connect" pos="-0.045 0.002 -0.085"/>
                            <inertial pos="0 0 0" quat="0.5 0.5 -0.5 0.5" mass="0.05" diaginertia="1e-5 1e-5 1e-5"/>
                            <joint name="left_foot_motor2_bar_ball_joint" type="ball" damping="0.005" frictionloss="0.01" stiffness="0" armature="0.001" limited="false"/>
                            <geom class="visual" mesh="left_shank_pitch_motor2_bar_link"/>
                        </body>
                    </body>
                    <body name="left_shank_pitch_motor1_link" pos="0 -0.011 -0.127">
                        <inertial pos="0 0 0" quat="0.5 0.5 -0.5 0.5" mass="0.05" diaginertia="1e-5 1e-5 1e-5"/>
                        <joint name="left_ankle_motor1_joint" class="left_foot_motor1"/>
                        <geom class="visual" mesh="left_shank_pitch_motor1_link"/>
                        <body name="left_shank_pitch_motor1_bar_link" pos="0.05 -0.013 -0.001">
                            <site name="left_foot_motor1_bar_connect" pos="0 0.024 -0.152"/>
                            <inertial pos="0 0 0" quat="0.5 0.5 -0.5 0.5" mass="0.05" diaginertia="1e-5 1e-5 1e-5"/>
                            <joint name="left_foot_motor1_bar_ball_joint" type="ball" damping="0.005" frictionloss="0.01" stiffness="0" armature="0.001" limited="false"/>
                            <geom class="visual" mesh="left_shank_pitch_motor1_bar_link"/>
                        </body>
                    </body>
                </body>
            </body>
      </body>
    </body>
    <body name="right_thigh_pitch_link" pos="0 -0.048521 -0.076847" quat="0.991445 0.130526 0 0">
        <inertial pos="0.0046972 -0.0370354 0.0002497" quat="0.394584 0.586067 -0.393152 0.588439" mass="1.3007" diaginertia="0.00238064 0.00207846 0.00131154"/>
        <joint name="right_hip_pitch_joint" class="hip_pitch"/>
        <geom class="visual" mesh="right_thigh_pitch_link"/>
        <body name="right_thigh_roll_link" pos="0 -0.074 0" quat="0.991445 -0.130526 0 0">
            <inertial pos="0.00184399 -4.33e-06 -0.0321936" quat="0.999702 -0.000215194 0.0223656 -0.00979228" mass="1.359" diaginertia="0.00302078 0.00288684 0.00125138"/>
            <joint name="right_hip_roll_joint" class="right_hip_roll"/>
            <geom class="visual" mesh="right_thigh_roll_link"/>
            <geom name="right_thigh" class="collision" type="capsule" size="0.07" fromto="0.0 0.01 0 0.0 0.01 -0.2"/>
            <body name="right_thigh_yaw_link" pos="0 0 -0.0875">
                <inertial pos="0.00216082 0.0097519 -0.0835711" quat="0.816219 0.0097719 -0.0399389 0.576277" mass="3.4517" diaginertia="0.0163407 0.0159406 0.00542163"/>
                <joint name="right_hip_yaw_joint" class="hip_yaw"/>
                <geom class="visual" mesh="right_thigh_yaw_link"/>
                <body name="right_shank_pitch_base_link" pos="0 0 -0.2031">
                    <inertial pos="0.00280042 -0.00416601 -0.135639" quat="0.392916 0.0123042 -0.00742701 0.919462" mass="2.3362" diaginertia="0.0138085 0.0137239 0.00176871"/>
                    <joint name="right_knee_pitch_joint" class="knee"/>
                    <geom class="visual" mesh="right_shank_pitch_base_link"/>
                    <geom name="right_shin" class="collision" type="capsule" size="0.07" fromto="0.02 0 0 0.02 0 -0.25"/>
                    <body name="right_foot_roll_link" pos="0 0 -0.28">
                        <inertial pos="0 -0.00180591 0" quat="0.5 0.5 -0.5 0.5" mass="0.11173" diaginertia="2.876e-05 2.445e-05 1.162e-05"/>
                        <joint name="right_ankle_roll_joint" axis="1 0 0" damping="0.5" frictionloss="1"/>
                        <geom class="visual" mesh="right_foot_roll_link"/>
                        <body name="right_foot_pitch_link">
                            <site name="right_foot" pos="0 0 0"/>
                            <site name="right_foot_motor1_connect" pos="0.05 0 0"/>
                            <site name="right_foot_motor2_connect" pos="0 -0.025 0"/>
                            <inertial pos="0.0239094 0.00023117 -0.0301994" quat="-0.00422828 0.703504 0.000787377 0.710678" mass="0.398" diaginertia="0.000900889 0.000843793 0.000145867"/>
                            <joint name="right_ankle_pitch_joint" axis="0 1 0" damping="0.5" frictionloss="1"/>
                            <geom class="visual" mesh="right_foot_pitch_link"/>
                            <geom name="right_foot" class="foot" pos="0.035 0 -0.035" size="0.10 0.04 0.008" type="box"/>
                        </body>
                    </body>
                    <body name="right_shank_pitch_motor2_link" pos="0.0205 0.0 -0.195">
                        <inertial pos="0 0 0" quat="0.5 0.5 -0.5 0.5" mass="0.05" diaginertia="1e-5 1e-5 1e-5"/>
                        <joint name="right_ankle_motor2_joint" class="foot_motor2"/>
                        <geom class="visual" mesh="right_shank_pitch_motor2_link"/>
                        <body name="right_shank_pitch_motor2_bar_link" pos="0.0245 -0.028 0">
                            <site name="right_foot_motor2_bar_connect" pos="-0.045 0.002 -0.085"/>
                            <inertial pos="0 0 0" quat="0.5 0.5 -0.5 0.5" mass="0.05" diaginertia="1e-5 1e-5 1e-5"/>
                            <joint name="right_foot_motor2_bar_ball_joint" type="ball" damping="0.005" frictionloss="0.01" stiffness="0" armature="0.001" limited="false"/>
                            <geom class="visual" mesh="right_shank_pitch_motor2_bar_link"/>
                        </body>
                    </body>
                    <body name="right_shank_pitch_motor1_link" pos="0 0.011 -0.127">
                        <inertial pos="0 0 0" quat="0.5 0.5 -0.5 0.5" mass="0.05" diaginertia="1e-5 1e-5 1e-5"/>
                        <joint name="right_ankle_motor1_joint" class="foot_motor1"/>
                        <geom class="visual" mesh="right_shank_pitch_motor1_link"/>
                        <body name="right_shank_pitch_motor1_bar_link" pos="0.05 0.013 -0.001">
                            <site name="right_foot_motor1_bar_connect" pos="0 -0.024 -0.152"/>
                            <inertial pos="0 0 0" quat="0.5 0.5 -0.5 0.5" mass="0.05" diaginertia="1e-5 1e-5 1e-5"/>
                            <joint name="right_foot_motor1_bar_ball_joint" type="ball" damping="0.005" frictionloss="0.01" stiffness="0" armature="0.001" limited="false"/>
                            <geom class="visual" mesh="right_shank_pitch_motor1_bar_link"/>
                        </body>
                    </body>
                </body>
            </body>
        </body>
    </body>
    </body>
  </worldbody>
  <contact>
    <!-- Feet floor. -->
    <pair name="left_foot_floor" geom1="left_foot" geom2="floor" condim="3" friction="1.0 1.0"/>
    <pair name="right_foot_floor" geom1="right_foot" geom2="floor" condim="3" friction="1.0 1.0"/>
    <!-- Hand floor -->
    <pair name="left_hand_floor" geom1="left_hand_collision" geom2="floor" condim="3" friction="0.8 0.8"/>
    <pair name="right_hand_floor" geom1="right_hand_collision" geom2="floor" condim="3" friction="0.8 0.8"/>

    <!-- body floor -->
    <pair name="body_floor" geom1="waist_yaw_link_collision" geom2="floor" condim="3" friction="0.8 0.8"/>

    <!-- elbow floor -->
    <pair name="left_elbow_floor" geom1="left_arm_collision" geom2="floor" condim="3" friction="0.8 0.8"/>
    <pair name="right_elbow_floor" geom1="right_arm_collision" geom2="floor" condim="3" friction="0.8 0.8"/>

    <!-- elbow body -->
    <pair name="left_elbow_body" geom1="left_arm_collision" geom2="waist_yaw_link_collision" condim="3" friction="0.8 0.8"/>
    <pair name="right_elbow_body" geom1="right_arm_collision" geom2="waist_yaw_link_collision" condim="3" friction="0.8 0.8"/>

    <!-- Feet / shin (opposite side) -->
    <pair name="left_foot_right_shin" geom1="left_foot" geom2="right_shin" condim="1"/>
    <pair name="right_foot_left_shin" geom1="right_foot" geom2="left_shin" condim="1"/>

    <!-- shin / floor -->
    <!-- <pair name="right_shin_floor" geom1="floor" geom2="right_shin" condim="1" friction="0.8 0.8"/>
    <pair name="left_shin_floor" geom1="floor" geom2="left_shin" condim="1" friction="0.8 0.8"/> -->

    <!-- Hand/thigh (same side) -->
    <pair name="left_hand_thigh" geom1="left_hand_collision" geom2="left_thigh" condim="1"/>
    <pair name="right_hand_thigh" geom1="right_hand_collision" geom2="right_thigh" condim="1"/>

    <!-- thigh / floor -->
    <pair name="left_thigh_floor" geom1="floor" geom2="left_thigh" condim="1" friction="0.8 0.8"/>
    <pair name="right_thigh_floor" geom1="floor" geom2="right_thigh" condim="1" friction="0.8 0.8"/>
    
    <!-- Left shin / Right shin. -->
    <pair name="left_shin_right_shin" geom1="left_shin" geom2="right_shin" condim="1"/>
  </contact>

<equality>
    <connect name="right_foot_motor1_connect"
             site1="right_foot_motor1_connect"
             site2="right_foot_motor1_bar_connect"
             solimp="0.9 0.95 0.001"
             solref="0.01 1"/>
    <connect name="right_foot_motor2_connect"
             site1="right_foot_motor2_connect"
             site2="right_foot_motor2_bar_connect"
             solimp="0.9 0.95 0.001"
             solref="0.01 1"/>
    <connect name="left_foot_motor1_connect"
             site1="left_foot_motor1_connect"
             site2="left_foot_motor1_bar_connect"
             solimp="0.9 0.95 0.001"
             solref="0.01 1"/>
    <connect name="left_foot_motor2_connect"
             site1="left_foot_motor2_connect"
             site2="left_foot_motor2_bar_connect"
             solimp="0.9 0.95 0.001"
             solref="0.01 1"/>
</equality>

  <actuator>
    <position class="waist_yaw" name="waist_yaw_joint" joint="waist_yaw_joint"/>

    <position class="shoulder_pitch" name="left_shoulder_pitch_joint" joint="left_shoulder_pitch_joint"/>
    <position class="left_shoulder_roll" name="left_shoulder_roll_joint" joint="left_shoulder_roll_joint"/>
    <position class="shoulder_yaw" name="left_shoulder_yaw_joint" joint="left_shoulder_yaw_joint"/>
    <position class="elbow" name="left_elbow_pitch_joint" joint="left_elbow_pitch_joint"/>
    <position class="wrist_yaw" name="left_wrist_yaw_joint" joint="left_wrist_yaw_joint"/>

    <position class="shoulder_pitch" name="right_shoulder_pitch_joint" joint="right_shoulder_pitch_joint"/>
    <position class="right_shoulder_roll" name="right_shoulder_roll_joint" joint="right_shoulder_roll_joint"/>
    <position class="shoulder_yaw" name="right_shoulder_yaw_joint" joint="right_shoulder_yaw_joint"/>
    <position class="elbow" name="right_elbow_pitch_joint" joint="right_elbow_pitch_joint"/>
    <position class="wrist_yaw" name="right_wrist_yaw_joint" joint="right_wrist_yaw_joint"/>
  


    <position class="hip_pitch" name="left_hip_pitch_joint" joint="left_hip_pitch_joint"/>
    <position class="left_hip_roll" name="left_hip_roll_joint" joint="left_hip_roll_joint"/>
    <position class="hip_yaw" name="left_hip_yaw_joint" joint="left_hip_yaw_joint"/>
    <position class="knee" name="left_knee_pitch_joint" joint="left_knee_pitch_joint"/>
    <position class="foot_motor2" name="left_ankle_motor2_joint" joint="left_ankle_motor2_joint"/>
    <position class="foot_motor1" name="left_ankle_motor1_joint" joint="left_ankle_motor1_joint"/>
    
    <position class="hip_pitch" name="right_hip_pitch_joint" joint="right_hip_pitch_joint"/>
    <position class="right_hip_roll" name="right_hip_roll_joint" joint="right_hip_roll_joint"/>
    <position class="hip_yaw" name="right_hip_yaw_joint" joint="right_hip_yaw_joint"/>
    <position class="knee" name="right_knee_pitch_joint" joint="right_knee_pitch_joint"/>
    <position class="foot_motor2" name="right_ankle_motor2_joint" joint="right_ankle_motor2_joint"/>
    <position class="foot_motor1" name="right_ankle_motor1_joint" joint="right_ankle_motor1_joint"/>
  </actuator>

  <sensor>
    <velocimeter site="imu" name="local_linvel_base_link"/>
    <gyro site="imu" name="gyro_base_link"/>
    <framezaxis objtype="site" objname="imu" name="upvector_base_link"/>
    <framexaxis objtype="site" objname="imu" name="forwardvector_base_link"/>
    <framequat objtype="site" objname="imu" name="orientation_base_link"/>
    <framelinvel objtype="site" objname="imu" name="global_linvel_base_link"/>
    <frameangvel objtype="site" objname="imu" name="global_angvel_base_link"/>
    <framelinvel objtype="site" objname="left_foot" name="global_linvel_left_foot"/>
    <framelinvel objtype="site" objname="right_foot" name="global_linvel_right_foot"/>
    <framezaxis objtype="site" objname="left_foot" name="upvector_left_foot"/>
    <framezaxis objtype="site" objname="right_foot" name="upvector_right_foot"/>
    <subtreecom name="robot_com" body="base_link"/>
    <subtreeangmom name="robot_mom" body="base_link"/>
  </sensor>
</mujoco>

and this is my scene file :

<mujoco model="grmini_t1 flat terrain scene">
  <include file="robot_mjx.xml"/>
  <!-- center 是相机或视野中心位置 extent 是场景最大边长（自动决定摄像头视角等） meansize 控制画面缩放/细节等默认估计尺寸 -->
  <statistic center="0 0 0.7" extent="1.2" meansize="0.04"/>

  <visual>
    <headlight diffuse=".8 .8 .8" ambient=".2 .2 .2" specular="1 1 1"/>
    <rgba force="1 0 0 1"/>
    <global azimuth="140" elevation="-20"/>
    <map force="0.01"/>
    <scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2"/>
    <quality shadowsize="8192"/>
  </visual>

  <asset>
    <texture name="grid" type="2d" builtin="checker" rgb1=".2 .3 .4" rgb2=".1 0.15 0.2" width="512" height="512"
      mark="cross" markrgb=".8 .8 .8"/>
    <material name="grid" texture="grid" texrepeat="1 1" texuniform="true"/>
  </asset>

  <worldbody>
    <geom name="floor" size="0 0 0.01" type="plane" material="grid"/>
    <body name="com_target" pos="0 0 0.5" mocap="true">
      <geom type="sphere" size="0.05" contype="0" conaffinity="0" rgba=".6 .3 .3 .2"/>
      <site type="sphere" size="0.02" rgba="0 0 1 1" group="1"/>
    </body>
    <body name="torso_target" pos="0 0 0.5" mocap="true">
      <!-- 把几何体从 sphere 改为 box，尺寸为 0.05m 的立方体 -->
      <geom type="box" size="0.05 0.05 0.05" contype="0" conaffinity="0" rgba=".6 .3 .3 .2"/>
      
      <!-- site 也改为 box，尺寸略小，用于可视化或控制目标 -->
      <site type="box" size="0.02 0.02 0.02" rgba="0 0 1 1" group="1"/>
    </body>
    <body name="left_foot_pitch_link_target" pos="0 0 0" quat="0 1 0 0" mocap="true">
      <geom type="sphere" size="0.05" contype="0" conaffinity="0" rgba=".6 .3 .3 .2"/>
      <site type="sphere" size="0.02" rgba="0 0 1 1" group="1"/>
    </body>
    <body name="right_foot_pitch_link_target" pos="0 0 0" quat="0 1 0 0" mocap="true">
      <geom type="sphere" size="0.05" contype="0" conaffinity="0" rgba=".6 .3 .3 .2"/>
      <site type="sphere" size="0.02" rgba="0 0 1 1" group="1"/>
    </body>
  </worldbody>

  <keyframe>
    <key name="home"
      qpos="
      0 0 0.710
      1 0 0 0
      0
      0 0 0 0 0 
      0 0 0 0 0 
      -0.309 0 0 0.65 0 -0.35 0 1 0 0 0 -0.35 1 0 0 0
      -0.309 0 0 0.65 0 -0.35 0 1 0 0 0 0.35 1 0 0 0
      "
      ctrl="
      0
      0 0 0 0 0 
      0 0 0 0 0 
      -0.309 0 0 0.65 0 -0.35
      -0.309 0 0 0.65 0 0.35
      "/>
    <key name="knees_bent"
      qpos="
      0 0 0.710
      1 0 0 0
      0
      0 0 0 0 0 
      0 0 0 0 0 
      -0.309 0 0 0.65 0 -0.35 0 1 0 0 0 -0.35 1 0 0 0
      -0.309 0 0 0.65 0 -0.35 0 1 0 0 0 0.35 1 0 0 0
      "
      ctrl="
      0
      0 0 0 0 0 
      0 0 0 0 0 
      -0.309 0 0 0.65 0 -0.35
      -0.309 0 0 0.65 0 0.35
    "/>
  </keyframe>
</mujoco>

sorry I can't upload my mesh files.

[kassasin]

I am using cuda to compute .

[kassasin]

And this is my script:

import traceback
from time import perf_counter

import jax
import jax.numpy as jnp
import mujoco as mj
import mujoco.mjx as mjx
import numpy as np
from jaxlie import SE3, SO3

from mjinx.components.barriers import JointBarrier
from mjinx.components.constraints import ModelEqualityConstraint
from mjinx.components.tasks import ComTask, FrameTask
from mjinx.configuration import integrate, update
from mjinx.problem import Problem
from mjinx.solvers import LocalIKSolver
from mjinx.visualize import BatchVisualizer

# === Mujoco ===

mj_model = mj.MjModel.from_xml_path("grmini_t1/scene_mjx.xml")
mjx_model = mjx.put_model(mj_model)
print(mjx_model.nq, mjx_model.nv)

q_min = mj_model.jnt_range[:, 0].copy()
q_max = mj_model.jnt_range[:, 1].copy()

# --- Mujoco visualization ---
# Initialize render window and launch it at the background
vis = BatchVisualizer("grmini_t1/scene_mjx.xml", n_models=5, alpha=0.2, record=True)

# === Mjinx ===
# --- Constructing the problem ---
# Creating problem formulation
problem = Problem(mjx_model, v_min=-1, v_max=1)

# Creating components of interest and adding them to the problem
joints_barrier = JointBarrier(
    "jnt_range",
    gain=0.1,
)
com_task = ComTask("com_task", cost=10.0, gain=50.0, mask=[1, 1, 1])
torso_task = FrameTask("torso_task", cost=1.0, gain=2.0, obj_name="base_link", mask=[0, 0, 0, 1, 1, 1])

# Feet (in stance)
left_foot_task = FrameTask(
    "left_foot_task",
    cost=20.0,
    gain=10.0,
    obj_name="left_foot_pitch_link",
    mask=[1, 1, 1, 1, 0, 1],
)
right_foot_task = FrameTask(
    "right_foot_task",
    cost=20.0,
    gain=10.0,
    obj_name="right_foot_pitch_link",
    mask=[1, 1, 1, 1, 0, 1],
)

model_equality_constraint = ModelEqualityConstraint()

problem.add_component(com_task)
problem.add_component(torso_task)
# TODO: fix this
problem.add_component(joints_barrier)
problem.add_component(left_foot_task)
problem.add_component(right_foot_task)
problem.add_component(model_equality_constraint)
# Initializing solver and its initial state
solver = LocalIKSolver(mjx_model, maxiter=10)

# Initializing initial condition
N_batch = 100
q0 = mj_model.keyframe("home").qpos
q = jnp.tile(q0, (N_batch, 1))

mjx_data = update(mjx_model, jnp.array(q0))

com0 = np.array(mjx_data.subtree_com[mjx_model.body_rootid[0]])
com_task.target_com = com0
# Get torso orientation and set it as target
torso_id = mjx.name2id(mjx_model, mj.mjtObj.mjOBJ_BODY, "base_link")
torso_quat = mjx_data.xquat[torso_id]
torso_task.target_frame = np.concatenate([np.zeros(3), torso_quat])

left_foot_id = mjx.name2id(mjx_model, mj.mjtObj.mjOBJ_BODY, "left_foot_pitch_link")
left_foot_pos = mjx_data.xpos[left_foot_id]
left_foot_quat = mjx_data.xquat[left_foot_id]
left_foot_task.target_frame = jnp.concatenate([left_foot_pos, left_foot_quat])

right_foot_id = mjx.name2id(mjx_model, mj.mjtObj.mjOBJ_BODY, "right_foot_pitch_link")
right_foot_pos = mjx_data.xpos[right_foot_id]
right_foot_quat = mjx_data.xquat[right_foot_id]
right_foot_task.target_frame = jnp.concatenate([right_foot_pos, right_foot_quat])

# Compiling the problem upon any parameters update
problem_data = problem.compile()
# --- Batching ---
print("Setting up batched computations...")
solver_data = jax.vmap(solver.init, in_axes=0)(v_init=jnp.zeros((N_batch, mjx_model.nv)))

with problem.set_vmap_dimension() as empty_problem_data:
    empty_problem_data.components["com_task"].target_com = 0

# Vmapping solve and integrate functions.
solve_jit = jax.jit(
    jax.vmap(
        solver.solve,
        in_axes=(0, 0, empty_problem_data),
    )
)
integrate_jit = jax.jit(jax.vmap(integrate, in_axes=(None, 0, 0, None)), static_argnames=["dt"])

# === Control loop ===
dt = 1e-2
ts = np.arange(0, 20, dt)

try:
    for t in ts:
        # Solving the instance of the problem
        # com_task.target_com = com0 - np.array([0, 0, 0.2 * np.sin(t) ** 2])
        com_task.target_com = np.array(
            [[0.0, 0.0, com0[2] - 0.3 * np.sin(t + 2 * np.pi * i / N_batch + np.pi / 2) ** 2] for i in range(N_batch)]
        )
        problem_data = problem.compile()
        opt_solution, solver_data = solve_jit(q, solver_data, problem_data)
        # Integrating
        q = integrate_jit(
            mjx_model,
            q,
            opt_solution.v_opt,
            dt,
        )
        # --- MuJoCo visualization ---
        indices = np.arange(0, N_batch, N_batch // vis.n_models)

        vis.update(q[:: N_batch // vis.n_models])

except KeyboardInterrupt:
    print("Finalizing the simulation as requested...")
except Exception:
    print(traceback.format_exc())
finally:
    if vis.record:
        vis.save_video(round(1 / dt))
    vis.close()

[kassasin]

I have modified equality_constraint.py

@jdc.pytree_dataclass
class JaxModelEqualityConstraint(JaxConstraint):
    """
    A JAX-based equality constraint derived from the simulation model.

    Equality constraints respresent all constrainst in the mujoco, which could be defined in
    <equality> tag in the xml file. More details could be found here: https://mujoco.readthedocs.io/en/stable/XMLreference.html#equality.

    This class utilized the fact, that during kinematics computations, the equality value are computed as well,
    and the jacobian of the equalities are stored in data.efc_J, and the residuals are stored in data.efc_pos.

    :param data: A MuJoCo simulation data structure containing model-specific constraint information.
    :return: A jax.numpy.ndarray with the positions corresponding to the equality constraints.
    """

    def __call__(self, data: mjx.Data) -> jnp.ndarray:
        """
        Compute the equality constraint values from the simulation data.

        This method selects the positions of equality constraints from the simulation data
        by filtering based on the constraint type and applying the mask indices.

        :param data: The MuJoCo simulation data.
        :return: A jax.numpy.ndarray containing the equality constraint values.
        """
        return data.efc_pos[jnp.where(data.efc_type == int(mj.mjtConstraint.mjCNSTR_EQUALITY),size=100, fill_value=-1)[0]][self.mask_idxs,]

    def compute_jacobian(self, data: mjx.Data) -> jnp.ndarray:
        """
        Compute the Jacobian matrix of the equality constraint.

        Retrieves the relevant rows of the Jacobian matrix from the simulation data,
        filtering by equality constraint type and applying the mask indices.

        :param data: The MuJoCo simulation data.
        :return: A jax.numpy.ndarray representing the Jacobian matrix of the equality constraints.
        """
        return data.efc_J[jnp.where(data.efc_type == int(mj.mjtConstraint.mjCNSTR_EQUALITY),size=100, fill_value=-1)[0], :][self.mask_idxs, :]

And it works on cpu, I will test on gpu later .

[domrachev03]

I am sincerely sorry for a late response.

I found an error that causes this problem, I would fix it with the next version of the mjinx. For now, you could try code from fix/eq_cnstr branch. I tried your example, and it launched on GPU, although I removed all the meshes, so I am not sure whether anything meaningful happened.

Please, let me know if it helped, and sorry once again for the delay

[kassasin]

I am sincerely sorry for a late response.

I found an error that causes this problem, I would fix it with the next version of the mjinx. For now, you could try code from fix/eq_cnstr branch. I tried your example, and it launched on GPU, although I removed all the meshes, so I am not sure whether anything meaningful happened.

Please, let me know if it helped, and sorry once again for the delay

Thanks! I will take a try!

[kassasin]

I am sincerely sorry for a late response.

I found an error that causes this problem, I would fix it with the next version of the mjinx. For now, you could try code from fix/eq_cnstr branch. I tried your example, and it launched on GPU, although I removed all the meshes, so I am not sure whether anything meaningful happened.

Please, let me know if it helped, and sorry once again for the delay

This method doesn't solve my problem, equality constraint is not totally concerned, even program runs well .
Sorry I cant send picture.

[domrachev03]

It is sad to hear, I am still working on fix for the second issue you raised. However, it is hard for me to figure out what exactly is wrong.

Is it possible to build a simplified model that has a similar problem? This would help me to debug the problem more precisely

Thank you!

[kassasin]

It is sad to hear, I am still working on fix for the second issue you raised. However, it is hard for me to figure out what exactly is wrong.

Is it possible to build a simplified model that has a similar problem? This would help me to debug the problem more precisely

Thank you!

Sorry, due to our work, we can't send file to outside . But for this problem, my solution mentioned above works will .

I have modified equality_constraint.py

@jdc.pytree_dataclass
class JaxModelEqualityConstraint(JaxConstraint):
    """
    A JAX-based equality constraint derived from the simulation model.

    Equality constraints respresent all constrainst in the mujoco, which could be defined in
    <equality> tag in the xml file. More details could be found here: https://mujoco.readthedocs.io/en/stable/XMLreference.html#equality.

    This class utilized the fact, that during kinematics computations, the equality value are computed as well,
    and the jacobian of the equalities are stored in data.efc_J, and the residuals are stored in data.efc_pos.

    :param data: A MuJoCo simulation data structure containing model-specific constraint information.
    :return: A jax.numpy.ndarray with the positions corresponding to the equality constraints.
    """

    def __call__(self, data: mjx.Data) -> jnp.ndarray:
        """
        Compute the equality constraint values from the simulation data.

        This method selects the positions of equality constraints from the simulation data
        by filtering based on the constraint type and applying the mask indices.

        :param data: The MuJoCo simulation data.
        :return: A jax.numpy.ndarray containing the equality constraint values.
        """
        return data.efc_pos[jnp.where(data.efc_type == int(mj.mjtConstraint.mjCNSTR_EQUALITY),size=100, fill_value=-1)[0]][self.mask_idxs,]

    def compute_jacobian(self, data: mjx.Data) -> jnp.ndarray:
        """
        Compute the Jacobian matrix of the equality constraint.

        Retrieves the relevant rows of the Jacobian matrix from the simulation data,
        filtering by equality constraint type and applying the mask indices.

        :param data: The MuJoCo simulation data.
        :return: A jax.numpy.ndarray representing the Jacobian matrix of the equality constraints.
        """
        return data.efc_J[jnp.where(data.efc_type == int(mj.mjtConstraint.mjCNSTR_EQUALITY),size=100, fill_value=-1)[0], :][self.mask_idxs, :]

And it works on cpu, I will test on gpu later .