chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rollout fix cosmetics

Browse Source
This commit is contained in:
Levi Burner
2025-02-21 08:44:17 -05:00
parent aca4e5d6ae
commit a0971ef749
2 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/rollout_test.py
+8 -4
View File
@@ -26,7 +26,6 @@ import numpy as np
import mujoco import mujoco
from mujoco import rollout from mujoco import rollout
# -------------------------- models used for testing --------------------------- # -------------------------- models used for testing ---------------------------
TEST_XML = r""" TEST_XML = r"""
@@ -885,12 +884,17 @@ class MuJoCoRolloutTest(parameterized.TestCase):
nthread = os.cpu_count() nthread = os.cpu_count()
nbatch = nthread nbatch = nthread
nstep = ((2**31) // (nstate*nbatch)) + 2 nstep = ((2**31) // (nstate * nbatch)) + 2
assert nstep * nstate * nbatch > 2**31 assert nstep * nstate * nbatch > 2**31
initial_state = np.random.randn(nbatch, nstate) initial_state = np.random.randn(nbatch, nstate)
rollout.rollout(model, [copy.copy(data) for _ in range(nthread)], rollout.rollout(
initial_state, nstep=nstep) model,
[copy.copy(data) for _ in range(nthread)],
initial_state,
nstep=nstep,
)
# -------------- Python implementation of rollout functionality ---------------- # -------------- Python implementation of rollout functionality ----------------