From 2019d42ca67d9eb7e3bb01eacb97ad98f9d0d4f7 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 22 Mar 2024 17:03:44 -0700 Subject: [PATCH] Small change to testspeed file. PiperOrigin-RevId: 618323795 Change-Id: I3a6307b8ae6d9af0d17c0ae85eae673d60112ddf --- mjx/mujoco/mjx/__init__.py | 1 + mjx/mujoco/mjx/_src/support.py | 1 - mjx/mujoco/mjx/_src/test_util.py | 82 ++++++++++++++++++++ mjx/mujoco/mjx/testspeed.py | 125 ++++++++++--------------------- 4 files changed, 124 insertions(+), 85 deletions(-) diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index e21b14fd..978f44ab 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -49,4 +49,5 @@ from mujoco.mjx._src.solver import solve from mujoco.mjx._src.support import full_m from mujoco.mjx._src.support import is_sparse from mujoco.mjx._src.support import mul_m +from mujoco.mjx._src.test_util import benchmark from mujoco.mjx._src.types import * diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 03966fa2..28314def 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -13,7 +13,6 @@ # limitations under the License. # ============================================================================== """Engine support functions.""" - from typing import Optional, Tuple, Union import jax diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py index b360fb63..36d07d9f 100644 --- a/mjx/mujoco/mjx/_src/test_util.py +++ b/mjx/mujoco/mjx/_src/test_util.py @@ -14,14 +14,96 @@ # ============================================================================== """Utilities for testing.""" +import os import sys +import time from typing import Dict, Tuple from xml.etree import ElementTree as ET from etils import epath +import jax import mujoco +# pylint: disable=g-importing-member +from mujoco.mjx._src import forward +from mujoco.mjx._src import io +# pylint: enable=g-importing-member import numpy as np + +def _measure(fn, *args) -> Tuple[float, float]: + """Reports jit time and op time for a function.""" + + beg = time.perf_counter() + compiled_fn = fn.lower(*args).compile() + end = time.perf_counter() + jit_time = end - beg + + beg = time.perf_counter() + result = compiled_fn(*args) + jax.block_until_ready(result) + end = time.perf_counter() + run_time = end - beg + + return jit_time, run_time + + +def benchmark( + m: mujoco.MjModel, + nstep: int = 1000, + batch_size: int = 1024, + unroll_steps: int = 1, + solver: str = 'cg', + iterations: int = 1, + ls_iterations: int = 4, +) -> Tuple[float, float, int]: + """Benchmark a model.""" + + xla_flags = os.environ.get('XLA_FLAGS', '') + xla_flags += ' --xla_gpu_triton_gemm_any=True' + os.environ['XLA_FLAGS'] = xla_flags + + m.opt.solver = { + 'cg': mujoco.mjtSolver.mjSOL_CG, + 'newton': mujoco.mjtSolver.mjSOL_NEWTON, + }[solver.lower()] + m.opt.iterations = iterations + m.opt.ls_iterations = ls_iterations + m = io.put_model(m) + + @jax.pmap + def init(key): + key = jax.random.split(key, batch_size // jax.device_count()) + + @jax.vmap + def random_init(key): + d = io.make_data(m) + qvel = 0.01 * jax.random.normal(key, shape=(m.nv,)) + d = d.replace(qvel=qvel) + return d + + return random_init(key) + + key = jax.random.split(jax.random.key(0), jax.device_count()) + d = init(key) + jax.block_until_ready(d) + + @jax.pmap + def unroll(d): + @jax.vmap + def step(d, _): + d = forward.step(m, d) + return d, None + + d, _ = jax.lax.scan(step, d, None, length=nstep, unroll=unroll_steps) + + return d + + jit_time, run_time = _measure(unroll, d) + steps = nstep * batch_size + + return jit_time, run_time, steps + + _ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity'] _DYN_TYPES = ['none', 'integrator', 'filter', 'filterexact'] _DYN_PRMS = ['0.189', '2.1'] diff --git a/mjx/mujoco/mjx/testspeed.py b/mjx/mujoco/mjx/testspeed.py index d66d475f..686e8cbc 100644 --- a/mjx/mujoco/mjx/testspeed.py +++ b/mjx/mujoco/mjx/testspeed.py @@ -14,110 +14,67 @@ # ============================================================================== """Run benchmarks on various devices.""" -import os -import time -from typing import Sequence, Tuple +from typing import Sequence from absl import app from absl import flags from etils import epath -import jax import mujoco from mujoco import mjx -FLAGS = flags.FLAGS - -flags.DEFINE_string('mjcf', None, 'path to model', required=True) -flags.DEFINE_integer('nstep', 1000, 'number of steps per rollout') -flags.DEFINE_integer('batch_size', 1024, 'number of parallel rollouts') -flags.DEFINE_integer('unroll', 1, 'loop unroll length') -flags.DEFINE_enum('solver', 'cg', ['cg', 'newton'], 'constraint solver') -flags.DEFINE_integer('iterations', 1, 'number of solver iterations') -flags.DEFINE_integer('ls_iterations', 4, 'number of linesearch iterations') -flags.DEFINE_enum('output', 'text', ['text', 'tsv'], 'format to print results') - - -def _measure(fn, *args) -> Tuple[float, float]: - """Reports jit time and op time for a function.""" - - beg = time.perf_counter() - compiled_fn = fn.lower(*args).compile() - end = time.perf_counter() - jit_time = end - beg - - beg = time.perf_counter() - result = compiled_fn(*args) - jax.block_until_ready(result) - end = time.perf_counter() - run_time = end - beg - - return jit_time, run_time +_MJCF = flags.DEFINE_string('mjcf', None, 'path to model', required=True) +_BASE_PATH = flags.DEFINE_string( + 'base_path', None, 'base path, defaults to mujoco.mjx resource path' +) +_NSTEP = flags.DEFINE_integer('nstep', 1000, 'number of steps per rollout') +_BATCH_SIZE = flags.DEFINE_integer( + 'batch_size', 1024, 'number of parallel rollouts' +) +_UNROLL = flags.DEFINE_integer('unroll', 1, 'loop unroll length') +_SOLVER = flags.DEFINE_enum( + 'solver', 'cg', ['cg', 'newton'], 'constraint solver' +) +_ITERATIONS = flags.DEFINE_integer( + 'iterations', 1, 'number of solver iterations' +) +_LS_ITERATIONS = flags.DEFINE_integer( + 'ls_iterations', 4, 'number of linesearch iterations' +) +_OUTPUT = flags.DEFINE_enum( + 'output', 'text', ['text', 'tsv'], 'format to print results' +) def _main(argv: Sequence[str]): - """Benchmark a model.""" - - xla_flags = os.environ.get('XLA_FLAGS', '') - xla_flags += ' --xla_gpu_triton_gemm_any=True' - os.environ['XLA_FLAGS'] = xla_flags - - f = epath.resource_path('mujoco.mjx') / 'test_data' / FLAGS.mjcf + """Runs testpeed function.""" + base_path = _BASE_PATH.value or epath.resource_path('mujoco.mjx') + f = base_path / 'test_data' / _MJCF.value m = mujoco.MjModel.from_xml_path(f.as_posix()) - m.opt.solver = { - 'cg': mujoco.mjtSolver.mjSOL_CG, - 'newton': mujoco.mjtSolver.mjSOL_NEWTON, - }[FLAGS.solver.lower()] - m.opt.iterations = FLAGS.iterations - m.opt.ls_iterations = FLAGS.ls_iterations - m = mjx.put_model(m) - if FLAGS.output == 'text': - print(f"Rolling out {FLAGS.nstep} 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( + m, + _NSTEP.value, + _BATCH_SIZE.value, + _UNROLL.value, + _SOLVER.value, + _ITERATIONS.value, + _LS_ITERATIONS.value, + ) - @jax.pmap - def init(key): - key = jax.random.split(key, FLAGS.batch_size // jax.device_count()) - - @jax.vmap - def random_init(key): - d = mjx.make_data(m) - qvel = 0.01 * jax.random.normal(key, shape=(m.nv,)) - d = d.replace(qvel=qvel) - return d - - return random_init(key) - - key = jax.random.split(jax.random.key(0), jax.device_count()) - d = init(key) - jax.block_until_ready(d) - - @jax.pmap - def unroll(d): - - @jax.vmap - def step(d, _): - d = mjx.step(m, d) - return d, None - - d, _ = jax.lax.scan(step, d, None, length=FLAGS.nstep, unroll=FLAGS.unroll) - - return d - - jit_time, run_time = _measure(unroll, d) - steps = FLAGS.nstep * FLAGS.batch_size - - if FLAGS.output == 'text': + name = argv[0] + if _OUTPUT.value == 'text': print(f""" -Summary for {FLAGS.batch_size} parallel rollouts +Summary for {_BATCH_SIZE.value} parallel rollouts Total JIT time: {jit_time:.2f} s Total simulation time: {run_time:.2f} s Total steps per second: { steps / run_time:.0f} Total realtime factor: { steps * m.opt.timestep / run_time:.2f} x Total time per step: { 1e6 * run_time / steps:.2f} µs""") - elif FLAGS.output == 'tsv': - name = argv[0].split('/')[-1].replace('testspeed_', '') - print(f"{name}\tjit: {jit_time:.2f}s\tsteps/second: {steps / run_time:.0f}") + elif _OUTPUT.value == 'tsv': + name = name.split('/')[-1].replace('testspeed_', '') + print(f'{name}\tjit: {jit_time:.2f}s\tsteps/second: {steps / run_time:.0f}') def main():