import mujoco_py
import timeit


T = 1000
NSIM = 100


def test():
    model = mujoco_py.load_model_from_path(
        "/Users/jonas/code/openai/pymj/xmls/claw.xml")
    pool = mujoco_py.MjSimPool([mujoco_py.MjSim(model)
                                for _ in range(NSIM)])
    for _ in range(T):
        pool.step()
        for sim in pool.sims:
            sim.data.qpos
            sim.data.qvel
            sim.data.ctrl
            sim.data.active_contacts_efc_pos

for i in range(5):
    NSIM = 1 * 10**i
    print(NSIM, timeit.timeit("test()", setup="from __main__ import test", number=1))