♻️ cleanup

scripts/train.py

@@ -8,8 +8,9 @@ _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)
-if sys.platform == "linux" and "DISPLAY" not in os.environ:
+# 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.
+if sys.platform == "linux":
     os.environ.setdefault("MUJOCO_GL", "osmesa")
 
 import hydra
@@ -73,9 +74,12 @@ def _init_clearml(choices: dict[str, str], remote: bool = False) -> Task:
         "registry.kube.optimize/worker-image:latest",
         docker_setup_bash_script=(
             "apt-get update && apt-get install -y --no-install-recommends "
-            "libosmesa6 libgl1-mesa-glx libglfw3 && rm -rf /var/lib/apt/lists/* "
-            "&& pip install 'jax[cuda12]' mujoco-mjx"
+            "libosmesa6-dev libgl1-mesa-glx libglfw3 && rm -rf /var/lib/apt/lists/* "
+            "&& pip install 'jax[cuda12]' mujoco-mjx PyOpenGL PyOpenGL-accelerate"
         ),
+        docker_arguments=[
+            "-e", "MUJOCO_GL=osmesa",
+        ],
     )
 
     req_file = pathlib.Path(hydra_utils.get_original_cwd()) / "requirements.txt"