[CI] Fix SOTA runs

ghstack-source-id: 425e5d7
Pull-Request: #3252

Author: vmoens

.github/unittest/linux_sota/scripts/run_all.sh

@@ -36,7 +36,6 @@ git config --global --add safe.directory '*'
 root_dir="$(git rev-parse --show-toplevel)"
 conda_dir="${root_dir}/conda"
 env_dir="${root_dir}/env"
-lib_dir="${env_dir}/lib"
 
 cd "${root_dir}"
 
@@ -72,14 +71,6 @@ fi
 printf "* Verified Python implementation: %s\n" "$python_impl"
 
 # 3. Install mujoco
-printf "* Installing mujoco and related\n"
-mkdir -p $root_dir/.mujoco
-cd $root_dir/.mujoco/
-#wget https://github.com/deepmind/mujoco/releases/download/2.1.1/mujoco-2.1.1-linux-x86_64.tar.gz
-#tar -xf mujoco-2.1.1-linux-x86_64.tar.gz
-wget https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz
-tar -xf mujoco210-linux-x86_64.tar.gz
-cd "${root_dir}"
 
 # 4. Install Conda dependencies
 printf "* Installing dependencies (except PyTorch)\n"
@@ -89,9 +80,6 @@ if ! grep -q "python=${PYTHON_VERSION}" "${this_dir}/environment.yml"; then
 fi
 cat "${this_dir}/environment.yml"
 
-export MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210
-#export MJLIB_PATH=$root_dir/.mujoco/mujoco-2.1.1/lib/libmujoco.so.2.1.1
-export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$root_dir/.mujoco/mujoco210/bin
 export SDL_VIDEODRIVER=dummy
 export MUJOCO_GL=egl
 export PYOPENGL_PLATFORM=egl
@@ -100,9 +88,7 @@ export COMPOSITE_LP_AGGREGATE=0
 
 conda env config vars set \
   MAX_IDLE_COUNT=1000 \
-  MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210 \
   DISPLAY=:99 \
-  LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$root_dir/.mujoco/mujoco210/bin \
   SDL_VIDEODRIVER=dummy \
   MUJOCO_GL=egl \
   PYOPENGL_PLATFORM=egl \
@@ -116,10 +102,6 @@ conda env update --file "${this_dir}/environment.yml" --prune
 conda deactivate
 conda activate "${env_dir}"
 
-# install d4rl
-pip install free-mujoco-py
-pip install git+https://github.com/Farama-Foundation/d4rl@master#egg=d4rl
-
 # TODO: move this down -- will break torchrl installation
 conda install -y -c conda-forge libstdcxx-ng=12
 ## find libstdc - search in the env's lib directory first, then fall back to conda packages
@@ -144,9 +126,6 @@ fi
 conda deactivate
 conda activate "${env_dir}"
 
-# compile mujoco-py (bc it's done at runtime for whatever reason someone thought it was a good idea)
-python -c """import gym;import d4rl"""
-
 # install ale-py: manylinux names are broken for CentOS so we need to manually download and
 # rename them
 
@@ -160,13 +139,11 @@ elif [[ ${#CU_VERSION} -eq 5 ]]; then
     CUDA_VERSION="${CU_VERSION:2:2}.${CU_VERSION:4:1}"
 fi
 echo "Using CUDA $CUDA_VERSION as determined by CU_VERSION ($CU_VERSION)"
-version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
-
 # submodules
 git submodule sync && git submodule update --init --recursive
 
 pip3 install ale-py -U
-pip3 install "gym[atari,accept-rom-license]" "gymnasium>=1.1.0" -U
+pip3 install "gymnasium[atari,accept-rom-license,mujoco]>=1.1.0" -U
 
 printf "Installing PyTorch with %s\n" "${CU_VERSION}"
 if [[ "$TORCH_VERSION" == "nightly" ]]; then

.github/unittest/linux_sota/scripts/test_sota.py

@@ -15,21 +15,6 @@
 ), "Composite LP must be set to False. Run this test with COMPOSITE_LP_AGGREGATE=0"
 
 commands = {
-    "dt": """python sota-implementations/decision_transformer/dt.py \
-  optim.pretrain_gradient_steps=55 \
-  optim.updates_per_episode=3 \
-  optim.warmup_steps=10 \
-  logger.backend= \
-  env.backend=gymnasium \
-  env.name=HalfCheetah-v4
-""",
-    "online_dt": """python sota-implementations/decision_transformer/online_dt.py \
-  optim.pretrain_gradient_steps=55 \
-  optim.updates_per_episode=3 \
-  optim.warmup_steps=10 \
-  env.backend=gymnasium \
-  logger.backend=
-""",
     "td3_bc": """python sota-implementations/td3_bc/td3_bc.py \
   optim.gradient_steps=55 \
   logger.backend=
@@ -39,7 +24,7 @@
   collector.frames_per_batch=20 \
   collector.num_workers=1 \
   logger.backend= \
-  env.backend=gym \
+  env.backend=gymnasium \
   logger.test_interval=10
 """,
     "ppo_mujoco": """python sota-implementations/ppo/ppo_mujoco.py \
@@ -57,7 +42,7 @@
   loss.mini_batch_size=20 \
   loss.ppo_epochs=2 \
   logger.backend= \
-  env.backend=gym \
+  env.backend=gymnasium \
   logger.test_interval=10
 """,
     "ddpg": """python sota-implementations/ddpg/ddpg.py \
@@ -84,7 +69,7 @@
   collector.frames_per_batch=20 \
   loss.mini_batch_size=20 \
   logger.backend= \
-  env.backend=gym \
+  env.backend=gymnasium \
   logger.test_interval=40
 """,
     "dqn_atari": """python sota-implementations/dqn/dqn_atari.py \
@@ -94,7 +79,7 @@
   buffer.batch_size=10 \
   loss.num_updates=1 \
   logger.backend= \
-  env.backend=gym \
+  env.backend=gymnasium \
   buffer.buffer_size=120
 """,
     "discrete_cql_online": """python sota-implementations/cql/discrete_cql_online.py \

.github/workflows/test-linux-sota.yml

@@ -26,14 +26,14 @@ jobs:
   tests:
     strategy:
       matrix:
-        python_version: ["3.9"]
-        cuda_arch_version: ["12.8"]
+        python_version: ["3.10"]
+        cuda_arch_version: ["13.0"]
       fail-fast: false
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       runner: linux.g5.4xlarge.nvidia.gpu
       repository: pytorch/rl
-      docker-image: "nvidia/cuda:12.2.0-devel-ubuntu22.04"
+      docker-image: "nvidia/cuda:13.0.2-cudnn-devel-ubuntu24.04"
       gpu-arch-type: cuda
       gpu-arch-version: ${{ matrix.cuda_arch_version }}
       timeout: 90