Small change to testspeed file.

PiperOrigin-RevId: 618323795
Change-Id: I3a6307b8ae6d9af0d17c0ae85eae673d60112ddf

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 *

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

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']

mjx/mujoco/mjx/testspeed.py

@@ -14,110 +14,67 @@
 # ==============================================================================
 """Run benchmarks on various devices."""
 
-import os
+from typing import Sequence
-import time
-from typing import Sequence, Tuple
 
 from absl import app
 from absl import flags
 from etils import epath
-import jax
 import mujoco
 from mujoco import mjx
 
-FLAGS = flags.FLAGS
+_MJCF = flags.DEFINE_string('mjcf', None, 'path to model', required=True)
-
+_BASE_PATH = flags.DEFINE_string(
-flags.DEFINE_string('mjcf', None, 'path to model', required=True)
+    'base_path', None, 'base path, defaults to mujoco.mjx resource path'
-flags.DEFINE_integer('nstep', 1000, 'number of steps per rollout')
+)
-flags.DEFINE_integer('batch_size', 1024, 'number of parallel rollouts')
+_NSTEP = flags.DEFINE_integer('nstep', 1000, 'number of steps per rollout')
-flags.DEFINE_integer('unroll', 1, 'loop unroll length')
+_BATCH_SIZE = flags.DEFINE_integer(
-flags.DEFINE_enum('solver', 'cg', ['cg', 'newton'], 'constraint solver')
+    'batch_size', 1024, 'number of parallel rollouts'
-flags.DEFINE_integer('iterations', 1, 'number of solver iterations')
+)
-flags.DEFINE_integer('ls_iterations', 4, 'number of linesearch iterations')
+_UNROLL = flags.DEFINE_integer('unroll', 1, 'loop unroll length')
-flags.DEFINE_enum('output', 'text', ['text', 'tsv'], 'format to print results')
+_SOLVER = flags.DEFINE_enum(
-
+    'solver', 'cg', ['cg', 'newton'], 'constraint solver'
-
+)
-def _measure(fn, *args) -> Tuple[float, float]:
+_ITERATIONS = flags.DEFINE_integer(
-  """Reports jit time and op time for a function."""
+    'iterations', 1, 'number of solver iterations'
-
+)
-  beg = time.perf_counter()
+_LS_ITERATIONS = flags.DEFINE_integer(
-  compiled_fn = fn.lower(*args).compile()
+    'ls_iterations', 4, 'number of linesearch iterations'
-  end = time.perf_counter()
+)
-  jit_time = end - beg
+_OUTPUT = flags.DEFINE_enum(
-
+    'output', 'text', ['text', 'tsv'], 'format to print results'
-  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 _main(argv: Sequence[str]):
-  """Benchmark a model."""
+  """Runs testpeed function."""
-
+  base_path = _BASE_PATH.value or epath.resource_path('mujoco.mjx')
-  xla_flags = os.environ.get('XLA_FLAGS', '')
+  f = base_path / 'test_data' / _MJCF.value
-  xla_flags += ' --xla_gpu_triton_gemm_any=True'
-  os.environ['XLA_FLAGS'] = xla_flags
-
-  f = epath.resource_path('mujoco.mjx') / 'test_data' / FLAGS.mjcf
   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 {_NSTEP.value} steps at dt = {m.opt.timestep:.3f}...')
-    print(f"Rolling out {FLAGS.nstep} 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
+  name = argv[0]
-  def init(key):
+  if _OUTPUT.value == 'text':
-    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':
     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':
+  elif _OUTPUT.value == 'tsv':
-    name = argv[0].split('/')[-1].replace('testspeed_', '')
+    name = name.split('/')[-1].replace('testspeed_', '')
-    print(f"{name}\tjit: {jit_time:.2f}s\tsteps/second: {steps / run_time:.0f}")
+    print(f'{name}\tjit: {jit_time:.2f}s\tsteps/second: {steps / run_time:.0f}')
 
 
 def main():