From dce38a75562e79ae9d1534e55948c61aeaa4b5fc Mon Sep 17 00:00:00 2001
From: Starrick Liu <73152103+StarrickLiu@users.noreply.github.com>
Date: Wed, 17 Jun 2026 20:07:03 +0800
Subject: [PATCH] Update Wall-X dependency and runtime setup (#105)

* Update Wall-X dependency and runtime setup

* Simplify FlashAttention installation notes
---
 README.md             |  3 ++-
 README.zh-CN.md       |  3 ++-
 requirements.txt      |  6 +++---
 scripts/run_libero.sh | 31 +++++++++++++++++++++++++++----
 4 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index c29cb56..dbbc33a 100644
--- a/README.md
+++ b/README.md
@@ -65,7 +65,8 @@ Install requirements:
 
 ```bash
 pip install -r requirements.txt
-MAX_JOBS=4 pip install flash-attn==2.7.4.post1 --no-build-isolation
+export FLASH_ATTN_CUDA_ARCHS=$(python -c 'import torch; print(f"{torch.cuda.get_device_capability()[0]}{torch.cuda.get_device_capability()[1]}")')
+MAX_JOBS=4 pip install flash-attn==2.8.3 --no-build-isolation
 ```
 
 Install DMuon, which is used by the default training configs:
diff --git a/README.zh-CN.md b/README.zh-CN.md
index 2902f05..81e73d9 100644
--- a/README.zh-CN.md
+++ b/README.zh-CN.md
@@ -29,7 +29,8 @@ conda activate wallx
 
 ```bash
 pip install -r requirements.txt
-MAX_JOBS=4 pip install flash-attn==2.7.4.post1 --no-build-isolation
+export FLASH_ATTN_CUDA_ARCHS=$(python -c 'import torch; print(f"{torch.cuda.get_device_capability()[0]}{torch.cuda.get_device_capability()[1]}")')
+MAX_JOBS=4 pip install flash-attn==2.8.3 --no-build-isolation
 ```
 
 安装默认训练配置使用的 DMuon：
diff --git a/requirements.txt b/requirements.txt
index e44b1a8..b7e5367 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
-torch==2.6.0
-torchvision==0.21.0
-torchaudio==2.6.0
+torch==2.10.0
+torchvision==0.25.0
+torchaudio==2.10.0
 transformers==5.2.0
 accelerate==1.13.0
 peft==0.18.1
diff --git a/scripts/run_libero.sh b/scripts/run_libero.sh
index c78b46f..6b83492 100755
--- a/scripts/run_libero.sh
+++ b/scripts/run_libero.sh
@@ -36,6 +36,19 @@ set -euo pipefail
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 SOURCE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
 
+usage() {
+    cat <<'EOF'
+Usage: bash scripts/run_libero.sh [CHECKPOINT_PATH]
+
+Environment:
+  CHECKPOINT_PATH=/path/to/checkpoint
+  CONFIG=/path/to/eval.yaml
+  LIBERO_PATH=/path/to/LIBERO
+  CUDA_ID=0
+  SMOKE=1
+EOF
+}
+
 if [[ -d "${SOURCE_ROOT}/third_party/LIBERO" ]]; then
     export PYTHONPATH="${SOURCE_ROOT}/third_party/LIBERO:${PYTHONPATH:-}"
 fi
@@ -61,15 +74,23 @@ DEFAULT_ALL_SUITES=(
     libero_10
 )
 
+case "${1:-}" in
+    -h|--help)
+        usage
+        exit 0
+        ;;
+esac
+
 if [[ $# -gt 1 ]]; then
-    echo "Usage: bash scripts/run_libero.sh [CHECKPOINT_PATH]" >&2
+    usage >&2
     exit 2
 fi
 if [[ $# -eq 1 ]]; then
     CHECKPOINT_PATH="$1"
 fi
 
-export CUDA_VISIBLE_DEVICES="${CUDA_ID:-0}"
+CUDA_IDS="${CUDA_ID:-0}"
+export CUDA_VISIBLE_DEVICES="${CUDA_IDS}"
 
 # MuJoCo offscreen rendering via NVIDIA EGL (required on headless GPU nodes).
 export MUJOCO_GL=egl
@@ -93,8 +114,10 @@ fi
 export __EGL_VENDOR_LIBRARY_FILENAMES="${EGL_VENDOR_FILE}"
 export LD_LIBRARY_PATH="/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib:/usr/local/nvidia/lib64:${LD_LIBRARY_PATH:-}"
 
-# After CUDA_VISIBLE_DEVICES remapping, MuJoCo only sees devices from index 0.
-export MUJOCO_EGL_DEVICE_ID="${MUJOCO_EGL_DEVICE_ID:-0}"
+# robosuite validates this value against CUDA_VISIBLE_DEVICES on common
+# headless EGL stacks, so default to the first requested physical GPU id.
+FIRST_CUDA_ID="${CUDA_IDS%%,*}"
+export MUJOCO_EGL_DEVICE_ID="${MUJOCO_EGL_DEVICE_ID:-${FIRST_CUDA_ID:-0}}"
 if [[ -d "${SOURCE_ROOT}/third_party/harrix/python" ]]; then
     export PYTHONPATH="${SOURCE_ROOT}/third_party/harrix/python:${PYTHONPATH:-}"
 fi