Update Wall-X dependency and runtime setup (#105)

* Update Wall-X dependency and runtime setup

* Simplify FlashAttention installation notes

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:

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：

requirements.txt

@@ -1,6 +1,6 @@
-torch==2.6.0
+torch==2.10.0
-torchvision==0.21.0
+torchvision==0.25.0
-torchaudio==2.6.0
+torchaudio==2.10.0
 transformers==5.2.0
 accelerate==1.13.0
 peft==0.18.1

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.
+# robosuite validates this value against CUDA_VISIBLE_DEVICES on common
-export MUJOCO_EGL_DEVICE_ID="${MUJOCO_EGL_DEVICE_ID:-0}"
+# 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