Data generation ends early

[Jiang-Yuxuan123]

Thanks for the great work!
I met a problem: when using the dataset_generation.sh script to generate data for the select_toy task, the task ends prematurely before completion, yet it is marked as successful. This may be caused by early simulator shutdown, rendering delays, or potentially incorrect task success criteria (e.g., the task may not actually require placing the toy in the box).
I used H100 with the command of MUJOCO_GL=egl bash dataset_generation.sh
A video demonstrating the issue is attached below:

demo_0_success_True.mp4

Any solutions to solve this issue?