diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 1e1101a5..e683633b 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -1486,10 +1486,11 @@ def _get_data_into_warp( if field.name in ( 'actuator_moment', 'contact', - 'efc_J', 'qM', 'qLD', 'qLDiagInv', + 'ten_J', + 'flexedge_J', ): continue if field.name.startswith('efc_'): diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index b47ccb03..8382a6cf 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -744,6 +744,22 @@ class DataIOTest(parameterized.TestCase): self.assertEqual(d_2.contact.frame.shape, (1, 9)) np.testing.assert_allclose(d_2.contact.frame, d.contact.frame) + def test_get_data_into_warp(self): + """Test get_data_into for impl='warp'.""" + + # TODO(taylorhowell): After put_data supports impl='warp' update test above + # and remove this test. + if not mjxw.WARP_INSTALLED: + self.skipTest('Warp is not installed.') + if not mjx_io.has_cuda_gpu_device(): + self.skipTest('No CUDA GPU device.') + + m = mujoco.MjModel.from_xml_string('') + d = mujoco.MjData(m) + mx = mjx.put_model(m, impl='warp') + dx = mjx.make_data(m, impl='warp') + mjx.get_data_into(d, mx, dx) + @parameterized.parameters('jax', 'c') def test_get_data_into_wrong_shape(self, impl): """Tests that get_data_into throwsif input and output shapes don't match."""