From b1941708ddb99c0469665b69f67024122b49035f Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Wed, 9 Oct 2024 07:14:02 -0700 Subject: [PATCH] Add some arena-allocated attributes to compare in mjData pickle/unpickle test. PiperOrigin-RevId: 684025320 Change-Id: Ic3282eaae74e43a33c2ed5f1b25ad38884487829 --- python/mujoco/bindings_test.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 32f5c0e6..6036366a 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -1365,7 +1365,21 @@ Euler integrator, semi-implicit in velocity. data2 = pickle.loads(pickle.dumps(self.data)) attr_to_compare = ( 'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos', - 'warning', 'energy' + 'warning', 'energy', 'contact', 'efc_J' + ) + self._assert_attributes_equal(data2, self.data, attr_to_compare) + for _ in range(10): + mujoco.mj_step(self.model, self.data) + mujoco.mj_step(self.model, data2) + self._assert_attributes_equal(data2, self.data, attr_to_compare) + + def test_pickle_mjdata_sparse(self): + self.model.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE + mujoco.mj_step(self.model, self.data) + data2 = pickle.loads(pickle.dumps(self.data)) + attr_to_compare = ( + 'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos', + 'warning', 'energy', 'contact', 'efc_J' ) self._assert_attributes_equal(data2, self.data, attr_to_compare) for _ in range(10):