From 0a60d555c69a715bb808b190aa5a3006c16faca9 Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Wed, 24 Apr 2024 10:29:39 -0700
Subject: [PATCH] Fix small bug in testspeed.

PiperOrigin-RevId: 627773755
Change-Id: Iebbcaa41cfdde126312e388d368843a4754abfb8
---
 mjx/mujoco/mjx/testspeed.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mjx/mujoco/mjx/testspeed.py b/mjx/mujoco/mjx/testspeed.py
index 686e8cbc..8065b3d5 100644
--- a/mjx/mujoco/mjx/testspeed.py
+++ b/mjx/mujoco/mjx/testspeed.py
@@ -47,8 +47,9 @@ _OUTPUT = flags.DEFINE_enum(
 
 def _main(argv: Sequence[str]):
   """Runs testpeed function."""
-  base_path = _BASE_PATH.value or epath.resource_path('mujoco.mjx')
-  f = base_path / 'test_data' / _MJCF.value
+  path = epath.resource_path('mujoco.mjx') / 'test_data'
+  path = _BASE_PATH.value or path
+  f = epath.Path(path) / _MJCF.value
   m = mujoco.MjModel.from_xml_path(f.as_posix())
 
   print(f'Rolling out {_NSTEP.value} steps at dt = {m.opt.timestep:.3f}...')