Can't run envpool mujoco environments

[rnilva]

Bug Description

Running envpool mujoco environments throws kernel errors.

System Information

• Ubuntu 22.04.4 LTS
• AMD Ryzen Threadripper 3970X 32-Core Processor

Steps to Reproduce

conda create -n arlbench-cpu python=3.10
conda activate arlbench-cpu
git clone https://github.com/automl/arlbench/
cd arlbench
make install-envpool
# This installs jaxlib==0.4.26 and envpool==0.8.4

cd examples
python run_arlbench.py -m environment=mujoco_hopper algorithm=sac search_space=sac_mujoco

# or 

python run_arlbench.py -m environment=mujoco_hopper algorithm=ppo search_space=ppo_mujoco

Expected Behavior

Run without errors.

Actual Behavior

1. Simulation instability warning appears:

   WARNING: Nan, Inf or huge value in CTRL at ACTUATOR 1/0. The simulation is unstable. Time = ...
   

2. Kernel errors occur:

   • "corrupted double-linked list"
   • "segmentation fault"

Additional Context

• Tested with fresh install on two different machines

• Running environment=procgen_bigfish_easy works correctly

• I know this could be the problem of envpool or its interaction with jax but just wonder if someone knows to solved this issue.