How to accelerate the sim speed when evaluate during training?

[Felixvillas]

Hi~ How to accelerate the model evaluation during the training phase? When I train a BC model using RGBD information, correspondingly, I need to obtain the RGBD information of the interactive environment during the evaluation phase. To speed up the acquisition of environmental RGBD information, I used MUJOCO_EGL="egl" during training. Similar to:

CUDA_VISIBLE_DEVICES=0 MUJOCO_EGL="egl" python robomimic/scripts/train.py --config robomimic/exps/templates/bc.json --dataset my_demo.hdf5

Unfortunately, I found that such a setting does not speed up the simulation speed of the environment during the evaluation phase of training.

However, when I just run the environment's test demo, MUJOCO_EGL="egl" can greatly accelerate the simulation speed of the environment. My test demo demo.py is as follows:

env = init_my_robosuite_env() ### init env, abbreviation
eps = 2
i_eps = 0
step_count = 0
start_time = time.time()
images = []
while i_eps < eps:
    obs = env.reset()
    images.append(obs["frontview_image"])
    d = False
    print(f"eps: {i_eps}, steps: {step_count}")
    ep_step_count = 0
    while not d:
        action = np.zeros(7)
        obs, r, d, _ = env.step(action)
        env._check_success()
        images.append(obs["frontview_image"])
        step_count += 1
        ep_step_count += 1
        print(f"eps: {i_eps}, steps: {step_count}, reward: {r}, d: {d}")
    i_eps += 1
end_time = time.time()

total_time = end_time - start_time
time_per_step = total_time / step_count
print(f"time_per_step: {time_per_step}")

# save video 
if video and save_video_flag:
    print(f"save video")
    save_video("./demo.mp4", [env.image_data])

The command I use to run this test demo is as follows:

CUDA_VISIBLE_DEVICES=0 MUJOCO_GL="egl" python demo.py

Notes: When I run demo.py, there is a Graphic process on the GPU, indicating that MUJOCO_GL="egl" is working. However, when I run train.py, there is no Graphic process on the GPU, indicating that MUJOCO_GL="egl" is not taking effect. I do not know why this situation is occurring.

Could you help me clarify this confusion? Thank you!

[Felixvillas]

@amandlek @cremebrule Can you help me answer the above questions? Thank you!

[amandlek]

This seems like it might be more related to mujoco and robosuite rather than robomimic. One other potential issue could be the policy forward pass time, which might be more of a bottleneck than simulation speed.