Fix bug in implementation of torquescale

PiperOrigin-RevId: 577466789
Change-Id: I3c8e5dd9b4af9a629e777f7ccad8417075b0046c

mjx/mujoco/mjx/_src/constraint.py

@@ -168,10 +168,10 @@ def _instantiate_weld(m: Model, d: Data) -> _Efc:
     jacdifr = 0.5 * jax.vmap(jac_fn)(jacdifr)
 
     j = jp.concatenate((jacdifp.T, jacdifr.T))
-    pos = jp.concatenate((cpos, crot))
+    pos = jp.concatenate((cpos, crot)).at[3:].mul(torquescale)
 
     # impedance, inverse constraint mass, reference acceleration
-    k, b, imp = _kbi(m, solref, solimp, math.norm(pos.at[3:].mul(torquescale)))
+    k, b, imp = _kbi(m, solref, solimp, math.norm(pos))
     invweight = m.body_invweight0[id1] + m.body_invweight0[id2]
     r = jp.maximum(invweight * (1 - imp) / imp, mujoco.mjMINVAL).repeat(3)
     aref = -b * (j @ d.qvel) - k * imp * pos

mjx/mujoco/mjx/_src/constraint_test.py

@@ -29,7 +29,7 @@ from mujoco.mjx._src.types import SolverType
 import numpy as np
 
 
-def _assert_eq(a, b, name, step, fname, atol=1e-3, rtol=1e-3):
+def _assert_eq(a, b, name, step, fname, atol=5e-3, rtol=5e-3):
   err_msg = f'mismatch: {name} at step {step} in {fname}'
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)