VictorMylle/RL-Sim-Framework
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

pin skrl/jax/mujoco/gymnasium versions; custom CUDA base image

Browse Source
This commit is contained in:
Victor Mylle
2026-06-10 09:05:48 +02:00
parent 56499ebe97
commit a6fbde798a
2 changed files with 7 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
scripts/train.py
View File

@@ -71,11 +71,11 @@ def _init_clearml(choices: dict[str, str], remote: bool = False) -> Task:
task = Task.init(project_name=project, task_name=task_name, tags=tags)
task.set_base_docker(
"registry.kube.optimize/worker-image:latest",
"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]' mujoco-mjx PyOpenGL PyOpenGL-accelerate"
"&& pip install 'jax[cuda12]==0.9.1' mujoco-mjx==3.5.0 PyOpenGL PyOpenGL-accelerate"
),
docker_arguments=[
"-e", "MUJOCO_GL=osmesa",