From aca4e5d6aef8c8c89cfe6a17d4289edc49dbc667 Mon Sep 17 00:00:00 2001 From: Levi Burner Date: Thu, 20 Feb 2025 21:12:52 -0500 Subject: [PATCH] rollout test for integer overflow in array size check --- python/mujoco/rollout_test.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/python/mujoco/rollout_test.py b/python/mujoco/rollout_test.py index f27dc02b..f37c2244 100644 --- a/python/mujoco/rollout_test.py +++ b/python/mujoco/rollout_test.py @@ -16,6 +16,7 @@ import concurrent.futures import copy +import os import threading from absl.testing import absltest @@ -875,6 +876,21 @@ class MuJoCoRolloutTest(parameterized.TestCase): model, [copy.copy(data) for i in range(3)], initial_state, control ) + @absltest.skip(reason='Takes a long time to run') + def test_large_state(self): + model = mujoco.MjModel.from_xml_string(TEST_XML) + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + data = mujoco.MjData(model) + + nthread = os.cpu_count() + nbatch = nthread + + nstep = ((2**31) // (nstate*nbatch)) + 2 + assert nstep * nstate * nbatch > 2**31 + + initial_state = np.random.randn(nbatch, nstate) + rollout.rollout(model, [copy.copy(data) for _ in range(nthread)], + initial_state, nstep=nstep) # -------------- Python implementation of rollout functionality ----------------