Error when using mujoco playground for vision-based environments

[Congying1112]

Dear Authors,
I have successfully run the MuJoCo playground and used it to train models without visual sensory input.
However, when I tried to run the script "train_vision_2.py" which is from the notebook "train_vision_ipynb", errors occurs as follows:

(mujoco_playground) (base) user@X299-UD4-Pro:/media/Data/workspace/mujoco_playground$ python learning/train_vision_2.py 
Using raytracer
Compiling /media/Data/workspace/madrona_mjx/external/madrona/src/mw/device/bvh.cpp
Compiling /media/Data/workspace/madrona_mjx/external/madrona/src/mw/device/memory.cpp
Compiling /media/Data/workspace/madrona_mjx/external/madrona/src/mw/device/host_print.cpp
Compiling /media/Data/workspace/madrona_mjx/external/madrona/src/mw/device/bvh_raycast.cpp
Compiling GPU engine code:
Initialization finished
Inferred batch not found, overriding manually
Inferred batch not found, overriding manually
Inferred batch not found, overriding manually
Inferred batch not found, overriding manually
E1020 16:34:25.092637  152068 pjrt_stream_executor_client.cc:3077] Execution of replica 0 failed: INTERNAL: There was an error before calling cuModuleGetFunction (1): cudaErrorInvalidValue : invalid argument
Traceback (most recent call last):
  File "/media/Data/workspace/mujoco_playground/learning/train_vision_2.py", line 127, in <module>
    tile(state.obs['pixels/view_0'][:64], 8)
         ~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/array.py", line 382, in __getitem__
    return indexing.rewriting_take(self, idx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/numpy/indexing.py", line 618, in rewriting_take
    if (result := _attempt_rewriting_take_via_slice(arr, idx, mode)) is not None:
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/numpy/indexing.py", line 594, in _attempt_rewriting_take_via_slice
    arr = lax.slice(
          ^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/lax/slicing.py", line 108, in slice
    return slice_p.bind(operand, start_indices=tuple(start_indices),
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/core.py", line 502, in bind
    return self._true_bind(*args, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/core.py", line 520, in _true_bind
    return self.bind_with_trace(prev_trace, args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/core.py", line 525, in bind_with_trace
    return trace.process_primitive(self, args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/core.py", line 1029, in process_primitive
    return primitive.impl(*args, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/lax/slicing.py", line 1414, in _slice_impl
    return dispatch.apply_primitive(dynamic_slice_p, x, *start_indices,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/dispatch.py", line 88, in apply_primitive
    outs = fun(*args)
           ^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/pjit.py", line 339, in cache_miss
    pgle_profiler) = _python_pjit_helper(fun, jit_info, *args, **kwargs)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/pjit.py", line 194, in _python_pjit_helper
    out_flat, compiled, profiler = _pjit_call_impl_python(*args_flat, **p.params)
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/pjit.py", line 1681, in _pjit_call_impl_python
    return compiled.unsafe_call(*args), compiled, pgle_profiler
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/profiler.py", line 334, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/media/Data/workspace/mujoco_playground/.venv/lib/python3.11/site-packages/jax/_src/interpreters/pxla.py", line 1297, in __call__
    results = self.xla_executable.execute_sharded(input_bufs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: There was an error before calling cuModuleGetFunction (1): cudaErrorInvalidValue : invalid argument
rt avg total time = 104.520607 ms

My OS is Ubuntu 22.04
Here are the results of nvidia-smi、nvcc and uv pip list

$ nvidia-smi
Mon Oct 20 17:47:19 2025       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 580.95.05              Driver Version: 580.95.05      CUDA Version: 13.0     |
+-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 3080        Off |   00000000:65:00.0  On |                  N/A |
|  0%   47C    P8             17W /  370W |     449MiB /  10240MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A            1242      G   /usr/lib/xorg/Xorg                      216MiB |
|    0   N/A  N/A            1491      G   /usr/bin/gnome-shell                     42MiB |
|    0   N/A  N/A          151800      G   /proc/self/exe                           60MiB |
|    0   N/A  N/A          151801      G   ...ersion=20251010-170048.933000         76MiB |
+-----------------------------------------------------------------------------------------+
$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2025 NVIDIA Corporation
Built on Wed_Jan_15_19:20:09_PST_2025
Cuda compilation tools, release 12.8, V12.8.61
Build cuda_12.8.r12.8/compiler.35404655_0
$ uv pip list
Package                   Version            Editable project location
------------------------- ------------------ --------------------------------------------
absl-py                   2.3.1
aiofiles                  25.1.0
annotated-types           0.7.0
anyio                     4.11.0
argon2-cffi               25.1.0
argon2-cffi-bindings      25.1.0
arrow                     1.3.0
astroid                   4.0.1
asttokens                 3.0.0
async-lru                 2.0.5
attrs                     25.4.0
babel                     2.17.0
beautifulsoup4            4.14.2
black                     24.10.0
bleach                    6.2.0
blinker                   1.9.0
brax                      0.13.0
certifi                   2025.10.5
cffi                      2.0.0
cfgv                      3.4.0
charset-normalizer        3.4.4
chex                      0.1.90
click                     8.3.0
cloudpickle               3.1.1
comm                      0.2.3
contourpy                 1.3.3
cupy                      13.6.0
cycler                    0.12.1
debugpy                   1.8.17
decorator                 5.2.1
defusedxml                0.7.1
dill                      0.4.0
distlib                   0.4.0
dm-control                1.0.34
dm-env                    1.6
dm-tree                   0.1.9
etils                     1.13.0
execnet                   2.1.1
executing                 2.2.1
fastjsonschema            2.21.2
fastrlock                 0.8.3
filelock                  3.20.0
flask                     3.1.2
flask-cors                6.0.1
flax                      0.10.6
fonttools                 4.60.1
fqdn                      1.5.1
fsspec                    2025.9.0
gitdb                     4.0.12
gitpython                 3.1.45
glfw                      2.10.0
h11                       0.16.0
httpcore                  1.0.9
httpx                     0.28.1
humanize                  4.13.0
identify                  2.6.15
idna                      3.11
immutabledict             4.2.2
importlab                 0.8.1
importlib-metadata        8.7.0
importlib-resources       6.5.2
iniconfig                 2.1.0
ipykernel                 7.0.1
ipython                   9.6.0
ipython-pygments-lexers   1.1.1
ipywidgets                8.1.7
isoduration               20.11.0
isort                     7.0.0
itsdangerous              2.2.0
jax                       0.5.3
jax-cuda12-pjrt           0.5.3
jax-cuda12-plugin         0.5.3
jaxlib                    0.5.3
jaxopt                    0.8.5
jedi                      0.19.2
jinja2                    3.1.6
json5                     0.12.1
jsonpointer               3.0.0
jsonschema                4.25.1
jsonschema-specifications 2025.9.1
jupyter                   1.1.1
jupyter-client            8.6.3
jupyter-console           6.6.3
jupyter-core              5.8.1
jupyter-events            0.12.0
jupyter-lsp               2.3.0
jupyter-server            2.17.0
jupyter-server-terminals  0.5.3
jupyterlab                4.4.9
jupyterlab-pygments       0.3.0
jupyterlab-server         2.27.3
jupyterlab-widgets        3.0.15
kiwisolver                1.4.9
labmaze                   1.0.6
lark                      1.3.0
libcst                    1.8.5
lxml                      6.0.2
madrona-mjx               0.0.1              /media/Data/workspace/madrona_mjx
markdown-it-py            4.0.0
markupsafe                3.0.3
matplotlib                3.10.7
matplotlib-inline         0.1.7
mccabe                    0.7.0
mdurl                     0.1.2
mediapy                   1.2.4
mistune                   3.1.4
ml-collections            1.1.0
ml-dtypes                 0.5.3
mpmath                    1.3.0
msgpack                   1.1.2
msgspec                   0.19.0
mujoco                    3.4.0.dev818719010
mujoco-mjx                3.3.7
mypy                      1.18.2
mypy-extensions           1.1.0
nbclient                  0.10.2
nbconvert                 7.16.6
nbformat                  5.10.4
nest-asyncio              1.6.0
networkx                  3.5
ninja                     1.13.0
nodeenv                   1.9.1
notebook                  7.4.7
notebook-shim             0.2.4
numpy                     2.3.3
nvidia-cublas-cu12        12.8.4.1
nvidia-cuda-cupti-cu12    12.8.90
nvidia-cuda-nvcc-cu12     12.9.86
nvidia-cuda-nvrtc-cu12    12.8.93
nvidia-cuda-runtime-cu12  12.8.90
nvidia-cudnn-cu12         9.10.2.21
nvidia-cufft-cu12         11.3.3.83
nvidia-cufile-cu12        1.13.1.3
nvidia-curand-cu12        10.3.9.90
nvidia-cusolver-cu12      11.7.3.90
nvidia-cusparse-cu12      12.5.8.93
nvidia-cusparselt-cu12    0.7.1
nvidia-nccl-cu12          2.27.3
nvidia-nvjitlink-cu12     12.8.93
nvidia-nvtx-cu12          12.8.90
onnx                      1.19.1
opt-einsum                3.4.0
optax                     0.2.6
orbax-checkpoint          0.11.24
orjson                    3.11.3
overrides                 7.7.0
packaging                 25.0
pandocfilters             1.5.1
parso                     0.8.5
pathspec                  0.12.1
pexpect                   4.9.0
pillow                    11.3.0
platformdirs              4.5.0
playground                0.0.5              /media/Data/workspace/mujoco_playground
pluggy                    1.6.0
pre-commit                4.3.0
prometheus-client         0.23.1
prompt-toolkit            3.0.52
protobuf                  6.32.1
psutil                    7.1.0
ptyprocess                0.7.0
pure-eval                 0.2.3
pycnite                   2024.7.31
pycparser                 2.23
pydantic                  2.12.2
pydantic-core             2.41.4
pydot                     4.0.1
pygments                  2.19.2
pyink                     24.10.1
pylint                    4.0.1
pyopengl                  3.1.10
pyparsing                 3.2.5
pytest                    8.4.2
pytest-xdist              3.8.0
python-dateutil           2.9.0.post0
python-json-logger        4.0.0
pytype                    2024.10.11
pyvers                    0.1.0
pyyaml                    6.0.3
pyzmq                     27.1.0
referencing               0.37.0
requests                  2.32.5
rfc3339-validator         0.1.4
rfc3986-validator         0.1.1
rfc3987-syntax            1.1.0
rich                      14.2.0
rpds-py                   0.27.1
rsl-rl-lib                3.1.1
ruff                      0.14.0
scipy                     1.16.2
send2trash                1.8.3
sentry-sdk                2.42.0
setuptools                80.9.0
simplejson                3.20.2
six                       1.17.0
smmap                     5.0.2
sniffio                   1.3.1
soupsieve                 2.8
stack-data                0.6.3
sympy                     1.14.0
tabulate                  0.9.0
tensorboardx              2.6.4
tensordict                0.10.0
tensorstore               0.1.78
terminado                 0.18.1
tinycss2                  1.4.0
toml                      0.10.2
tomlkit                   0.13.3
toolz                     1.0.0
torch                     2.8.0
torchvision               0.23.0
tornado                   6.5.2
tqdm                      4.67.1
traitlets                 5.14.3
treescope                 0.1.10
trimesh                   4.8.3
triton                    3.4.0
types-python-dateutil     2.9.0.20251008
typing-extensions         4.15.0
typing-inspection         0.4.2
uri-template              1.3.0
urllib3                   2.5.0
virtualenv                20.35.3
vulkan                    1.3.275.1
wandb                     0.22.2
warp-lang                 1.10.0.dev20251014
wcwidth                   0.2.14
webcolors                 24.11.1
webencodings              0.5.1
websocket-client          1.9.0
werkzeug                  3.1.3
widgetsnbextension        4.0.14
wrapt                     1.17.3
zipp                      3.23.0

Thank you so much.