How to install pytorch with mujoco-playground? How to run train_rsl_rl.py?

I just want to run the train_rsl_rl.py, and use pytorch with mujoco-playground. But I face to a problem about installing pytorch.
I use the following instruction to construct my virtual environment:
```
pip install -U "jax[cuda12]"
pip install -e ".[all]"
pip install torch torchvision torchaudio
```
but when I run the script, an error occurs:
```
E0614 09:34:18.652694 35835 cuda_dnn.cc:522] Loaded runtime CuDNN library: 9.5.1 but source was compiled with: 9.8.0. CuDNN library needs to have matching major version and equal or higher minor version. If using a binary install, upgrade your CuDNN library. If building from sources, make sure the library loaded at runtime is compatible with the version specified during compile configuration.
E0614 09:34:18.653664 35835 cuda_dnn.cc:522] Loaded runtime CuDNN library: 9.5.1 but source was compiled with: 9.8.0. CuDNN library needs to have matching major version and equal or higher minor version. If using a binary install, upgrade your CuDNN library. If building from sources, make sure the library loaded at runtime is compatible with the version specified during compile configuration.
E0614 09:34:18.712497 35835 cuda_dnn.cc:522] Loaded runtime CuDNN library: 9.5.1 but source was compiled with: 9.8.0. CuDNN library needs to have matching major version and equal or higher minor version. If using a binary install, upgrade your CuDNN library. If building from sources, make sure the library loaded at runtime is compatible with the version specified during compile configuration.
E0614 09:34:18.719670 35835 cuda_dnn.cc:522] Loaded runtime CuDNN library: 9.5.1 but source was compiled with: 9.8.0. CuDNN library needs to have matching major version and equal or higher minor version. If using a binary install, upgrade your CuDNN library. If building from sources, make sure the library loaded at runtime is compatible with the version specified during compile configuration.
Traceback (most recent call last):
File "/home/vincentfei/humanoid/mujoco/mujoco_playground/learning/train_jax_ppo.py", line 506, in
app.run(main)
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/absl/app.py", line 316, in run
_run_main(main, args)
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/absl/app.py", line 261, in _run_main
sys.exit(main(argv))
^^^^^^^^^^
File "/home/vincentfei/humanoid/mujoco/mujoco_playground/learning/train_jax_ppo.py", line 260, in main
env = registry.load(_ENV_NAME.value, config=env_cfg)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/humanoid/mujoco/mujoco_playground/mujoco_playground/_src/registry.py", line 59, in load
return dm_control_suite.load(env_name, config, config_overrides)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/humanoid/mujoco/mujoco_playground/mujoco_playground/_src/dm_control_suite/__init__.py", line 157, in load
return _envs[env_name](config=config, config_overrides=config_overrides)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/humanoid/mujoco/mujoco_playground/mujoco_playground/_src/dm_control_suite/cartpole.py", line 98, in __init__
self._mjx_model = mjx.put_model(self._mj_model)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/mujoco/mjx/_src/io.py", line 217, in put_model
return _strip_weak_type(model)
^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/mujoco/mjx/_src/io.py", line 38, in _strip_weak_type
return jax.tree_util.tree_map(f, tree)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in tree_map
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/tree_util.py", line 361, in
return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/mujoco/mjx/_src/io.py", line 35, in f
return leaf.astype(jax.dtypes.canonicalize_dtype(leaf.dtype))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py", line 122, in _astype
return lax_numpy.astype(self, dtype, copy=copy, device=device)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py", line 5647, in astype
result = lax_internal._convert_element_type(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 1727, in _convert_element_type
return convert_element_type_p.bind(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/core.py", line 531, in bind
return self._true_bind(*args, **params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/core.py", line 551, in _true_bind
return self.bind_with_trace(prev_trace, args, params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 4897, in _convert_element_type_bind_with_trace
operand = core.Primitive.bind_with_trace(convert_element_type_p, trace, args, params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/core.py", line 556, in bind_with_trace
return trace.process_primitive(self, args, params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/core.py", line 1060, in process_primitive
return primitive.impl(*args, **params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/vincentfei/miniconda3/envs/mujoco-playground-test/lib/python3.11/site-packages/jax/_src/dispatch.py", line 88, in apply_primitive
outs = fun(*args)
^^^^^^^^^^
jaxlib._jax.XlaRuntimeError: FAILED_PRECONDITION: DNN library initialization failed. Look at the errors above for more details.
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
```

This error is related to the version of CuDNN. I will give the version of CuDNN before and after I install pytorch.

Before I install the pytorch
```
nvidia-cublas-cu12 12.9.1.4
nvidia-cuda-cupti-cu12 12.9.79
nvidia-cuda-nvcc-cu12 12.9.86
nvidia-cuda-runtime-cu12 12.9.79
nvidia-cudnn-cu12 9.10.2.21
nvidia-cufft-cu12 11.4.1.4
nvidia-cusolver-cu12 11.7.5.82
nvidia-cusparse-cu12 12.5.10.65
nvidia-nccl-cu12 2.27.3
nvidia-nvjitlink-cu12 12.9.86
nvidia-nvshmem-cu12 3.2.5
```

After I install the pytorch
```
nvidia-cublas-cu12 12.6.4.1
nvidia-cuda-cupti-cu12 12.6.80
nvidia-cuda-nvcc-cu12 12.9.86
nvidia-cuda-nvrtc-cu12 12.6.77
nvidia-cuda-runtime-cu12 12.6.77
nvidia-cudnn-cu12 9.5.1.17
nvidia-cufft-cu12 11.3.0.4
nvidia-cufile-cu12 1.11.1.6
nvidia-curand-cu12 10.3.7.77
nvidia-cusolver-cu12 11.7.1.2
nvidia-cusparse-cu12 12.5.4.2
nvidia-cusparselt-cu12 0.6.3
nvidia-nccl-cu12 2.26.2
nvidia-nvjitlink-cu12 12.6.85
nvidia-nvshmem-cu12 3.2.5
nvidia-nvtx-cu12 12.6.77
```

After install the pytorch, the nvidia package is degraded. How to solve this problem? Is there anybody can give me a methods to install the pytorch environment?