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

10
requirements.txt
View File

@@ -1,11 +1,11 @@
torch torch
gymnasium gymnasium==1.2.3
hydra-core hydra-core
omegaconf omegaconf
mujoco mujoco==3.5.0
mujoco-mjx mujoco-mjx==3.5.0
jax jax==0.9.1
skrl[torch] skrl[torch]==1.4.3
clearml clearml
imageio imageio
imageio-ffmpeg imageio-ffmpeg

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 = Task.init(project_name=project, task_name=task_name, tags=tags)
task.set_base_docker( task.set_base_docker(
"registry.kube.optimize/worker-image:latest", "git.victormylle.be/victormylle/simple-rl-framework:latest",
docker_setup_bash_script=( docker_setup_bash_script=(
"apt-get update && apt-get install -y --no-install-recommends " "apt-get update && apt-get install -y --no-install-recommends "
"libosmesa6-dev libgl1-mesa-glx libglfw3 && rm -rf /var/lib/apt/lists/* " "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=[ docker_arguments=[
"-e", "MUJOCO_GL=osmesa", "-e", "MUJOCO_GL=osmesa",