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

Merge pull request #2078 from hartikainen:enable-mjb-for-mjx-testspeed

Browse Source 
PiperOrigin-RevId: 682328610
Change-Id: I1329d8f7a0a47234e3eda61c457567825c89bd88
This commit is contained in:
Copybara-Service
2024-10-04 08:22:53 -07:00
parent 42c1718d57 6f031f2293
commit 3ea8be2dcb
2 changed files with 9 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/testspeed.py
+5 -2
View File
@@ -22,7 +22,7 @@ from etils import epath
import mujoco import mujoco
from mujoco import mjx from mujoco import mjx
_MJCF = flags.DEFINE_string('mjcf', None, 'path to model', required=True) _MJCF = flags.DEFINE_string('mjcf', None, 'path to model `.xml` or `.mjb`', required=True)
_BASE_PATH = flags.DEFINE_string( _BASE_PATH = flags.DEFINE_string(
'base_path', None, 'base path, defaults to mujoco.mjx resource path' 'base_path', None, 'base path, defaults to mujoco.mjx resource path'
) )
@@ -50,7 +50,10 @@ def _main(argv: Sequence[str]):
path = epath.resource_path('mujoco.mjx') / 'test_data' path = epath.resource_path('mujoco.mjx') / 'test_data'
path = _BASE_PATH.value or path path = _BASE_PATH.value or path
f = epath.Path(path) / _MJCF.value f = epath.Path(path) / _MJCF.value
m = mujoco.MjModel.from_xml_path(f.as_posix()) if f.suffix == '.mjb':
m = mujoco.MjModel.from_binary_path(f.as_posix())
else:
m = mujoco.MjModel.from_xml_path(f.as_posix())
print(f'Rolling out {_NSTEP.value} steps at dt = {m.opt.timestep:.3f}...') print(f'Rolling out {_NSTEP.value} steps at dt = {m.opt.timestep:.3f}...')
jit_time, run_time, steps = mjx.benchmark( jit_time, run_time, steps = mjx.benchmark(
mjx/mujoco/mjx/viewer.py
+4 -1
View File
@@ -51,7 +51,10 @@ def _main(argv: Sequence[str]) -> None:
jax.config.update('jax_debug_nans', True) jax.config.update('jax_debug_nans', True)
print(f'Loading model from: {_MODEL_PATH.value}.') print(f'Loading model from: {_MODEL_PATH.value}.')
m = mujoco.MjModel.from_xml_path(_MODEL_PATH.value) if _MODEL_PATH.value.endswith('.mjb'):
m = mujoco.MjModel.from_binary_path(_MODEL_PATH.value)
else:
m = mujoco.MjModel.from_xml_path(_MODEL_PATH.value)
d = mujoco.MjData(m) d = mujoco.MjData(m)
mx = mjx.put_model(m) mx = mjx.put_model(m)
dx = mjx.put_data(m, d) dx = mjx.put_data(m, d)