jax[cuda12] on linux for GPU; EGL headless render; non-fatal video

scripts/train.py

@@ -8,10 +8,12 @@ _PROJECT_ROOT = str(pathlib.Path(__file__).resolve().parent.parent)
 if _PROJECT_ROOT not in sys.path:
     sys.path.insert(0, _PROJECT_ROOT)
 
-# Headless rendering: use OSMesa on Linux servers (must be set before mujoco import).
-# Always default on Linux — Docker containers may have DISPLAY set without a real X server.
+# Headless rendering on Linux servers (must be set before mujoco import).
+# EGL renders on the GPU directly (right for NVIDIA nodes) and avoids the
+# brittle OSMesa/PyOpenGL stack. Forced (not setdefault) so a stale
+# `-e MUJOCO_GL=osmesa` baked into a remote task can't override it.
 if sys.platform == "linux":
-    os.environ.setdefault("MUJOCO_GL", "osmesa")
+    os.environ["MUJOCO_GL"] = "egl"
 
 import hydra
 import hydra.utils as hydra_utils
@@ -74,11 +76,11 @@ def _init_clearml(choices: dict[str, str], remote: bool = False) -> Task:
         "git.victormylle.be/victormylle/simple-rl-framework:latest",
         docker_setup_bash_script=(
             "apt-get update && apt-get install -y --no-install-recommends "
-            "libosmesa6-dev libgl1-mesa-glx libglfw3 && rm -rf /var/lib/apt/lists/* "
-            "&& pip install 'jax[cuda12]==0.9.1' mujoco-mjx==3.5.0 PyOpenGL PyOpenGL-accelerate"
+            "libegl1 libgl1 libglfw3 libosmesa6 && rm -rf /var/lib/apt/lists/* "
+            "&& pip install 'jax[cuda12]==0.9.1' mujoco-mjx==3.5.0"
         ),
         docker_arguments=[
-            "-e", "MUJOCO_GL=osmesa",
+            "-e", "MUJOCO_GL=egl",
         ],
     )