Skip to content

Commit ed66a0f

Browse files
vmoensvmoens
committed
[Feature] Add LIBERO environment adapter
ghstack-source-id: 5e155e7 Pull-Request: #3877
1 parent 025c8d2 commit ed66a0f

14 files changed

Lines changed: 1411 additions & 0 deletions

File tree

‎.github/labeler.yml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,10 @@
3737
- changed-files:
3838
- any-glob-to-any-file: ['torchrl/envs/libs/jumanji.py']
3939

40+
"Environments/libero":
41+
- changed-files:
42+
- any-glob-to-any-file: ['torchrl/envs/libs/libero.py', 'test/libs/test_libero.py']
43+
4044
"Environments/meltingpot":
4145
- changed-files:
4246
- any-glob-to-any-file: ['torchrl/envs/libs/meltingpot.py']

‎.github/unittest/linux_libs/scripts_libero/install.sh‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
#!/usr/bin/env bash
2+
3+
set -euxo pipefail
4+
5+
unset PYTORCH_VERSION
6+
7+
root_dir="$(git rev-parse --show-toplevel)"
8+
env_dir="${root_dir}/env"
9+
source "${env_dir}/bin/activate"
10+
export PATH="$HOME/.local/bin:$PATH"
11+
12+
uv_pip_install() {
13+
uv pip install --no-progress --python "${env_dir}/bin/python" "$@"
14+
}
15+
16+
if [ "${CU_VERSION:-}" == cpu ] ; then
17+
torch_index="https://download.pytorch.org/whl/nightly/cpu"
18+
stable_torch_index="https://download.pytorch.org/whl/cpu"
19+
else
20+
torch_index="https://download.pytorch.org/whl/nightly/${CU_VERSION}"
21+
stable_torch_index="https://download.pytorch.org/whl/${CU_VERSION}"
22+
fi
23+
24+
git submodule sync && git submodule update --init --recursive
25+
26+
uv_pip_install \
27+
cloudpickle \
28+
coverage \
29+
expecttest \
30+
future \
31+
hydra-core \
32+
hypothesis \
33+
importlib_metadata \
34+
orjson \
35+
packaging \
36+
psutil \
37+
"pyvers>=0.2.0,<0.3.0" \
38+
pybind11[global] \
39+
pytest \
40+
pytest-asyncio \
41+
pytest-cov \
42+
pytest-error-for-skips \
43+
pytest-instafail \
44+
pytest-json-report \
45+
pytest-mock \
46+
pytest-rerunfailures \
47+
pytest-timeout \
48+
pyyaml \
49+
scipy \
50+
setuptools \
51+
wheel
52+
53+
if [[ "$TORCH_VERSION" == "nightly" ]]; then
54+
uv_pip_install --upgrade --pre torch --index-url "${torch_index}"
55+
elif [[ "$TORCH_VERSION" == "stable" ]]; then
56+
uv_pip_install --upgrade torch --index-url "${stable_torch_index}"
57+
else
58+
echo "Failed to install pytorch"
59+
exit 1
60+
fi
61+
62+
# Install TensorDict before TorchRL. Nightly CI validates against TensorDict main.
63+
if [[ "$RELEASE" == 0 ]]; then
64+
uv_pip_install --no-build-isolation --no-deps git+https://github.com/pytorch/tensordict.git
65+
else
66+
uv_pip_install --no-deps tensordict
67+
fi
68+
69+
uv_pip_install -e . --no-build-isolation --no-deps
70+
71+
# LIBERO is source-only. Its root is a namespace-package parent; keep it on
72+
# PYTHONPATH at test time instead of relying on the upstream editable install.
73+
libero_dir="${root_dir}/libero-src"
74+
rm -rf "${libero_dir}"
75+
git clone --depth 1 https://github.com/Lifelong-Robot-Learning/LIBERO.git "${libero_dir}"
76+
77+
uv_pip_install \
78+
"bddl==1.0.1" \
79+
easydict \
80+
"gym==0.25.2" \
81+
h5py \
82+
imageio \
83+
matplotlib \
84+
"numpy<2" \
85+
opencv-python \
86+
"robosuite==1.4.0" \
87+
termcolor \
88+
tqdm
89+
90+
timeout 120s python -c "import functorch; import tensordict; import torchrl"
91+
92+
export LIBERO_CONFIG_PATH="${root_dir}/.libero-ci"
93+
timeout 120s env PYTHONPATH="${libero_dir}:${PYTHONPATH:-}" python -c "from torchrl.envs.libs.libero import _ensure_libero_config; _ensure_libero_config()"
94+
timeout 120s env PYTHONPATH="${libero_dir}:${PYTHONPATH:-}" python -c "from libero.libero import benchmark; from libero.libero.envs import OffScreenRenderEnv; print(sorted(benchmark.get_benchmark_dict()))"

‎.github/unittest/linux_libs/scripts_libero/post_process.sh‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
#!/usr/bin/env bash
2+
3+
set -e
4+
5+
source ./env/bin/activate

‎.github/unittest/linux_libs/scripts_libero/run_all.sh‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
#!/usr/bin/env bash
2+
3+
set -euxo pipefail
4+
5+
this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
6+
7+
bash "${this_dir}/setup_env.sh"
8+
bash "${this_dir}/install.sh"
9+
PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_version.sh" "$TORCH_VERSION"
10+
PYTHON=./env/bin/python bash "$(git rev-parse --show-toplevel)/.github/unittest/helpers/assert_torch_tensordict_versions.sh" "$TORCH_VERSION"
11+
bash "${this_dir}/run_test.sh"
12+
bash "${this_dir}/post_process.sh"

‎.github/unittest/linux_libs/scripts_libero/run_test.sh‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
#!/usr/bin/env bash
2+
3+
set -euxo pipefail
4+
5+
root_dir="$(git rev-parse --show-toplevel)"
6+
env_dir="${root_dir}/env"
7+
libero_dir="${root_dir}/libero-src"
8+
source "${env_dir}/bin/activate"
9+
10+
export PYTHONPATH="${libero_dir}:${PYTHONPATH:-}"
11+
export LIBERO_CONFIG_PATH="${root_dir}/.libero-ci"
12+
export PYTORCH_TEST_WITH_SLOW='1'
13+
export LAZY_LEGACY_OP=False
14+
export MUJOCO_GL=egl
15+
export PYOPENGL_PLATFORM=egl
16+
export SDL_VIDEODRIVER=dummy
17+
export DISPLAY=:99
18+
export MKL_THREADING_LAYER=GNU
19+
export CUDA_VISIBLE_DEVICES=0
20+
21+
timeout 300s python -m torch.utils.collect_env
22+
git config --global --add safe.directory '*'
23+
24+
Xvfb :99 -screen 0 1024x768x24 &
25+
26+
timeout 120s python -c 'import torch; t = torch.ones([2, 2], device="cuda:0" if torch.cuda.is_available() else "cpu"); print(t); print("tensor device:" + str(t.device))'
27+
timeout 120s python -c 'from torchrl.envs.libs.libero import _ensure_libero_config; _ensure_libero_config()'
28+
timeout 120s python -c 'from libero.libero import benchmark; from libero.libero.envs import OffScreenRenderEnv; print("LIBERO suites:", sorted(benchmark.get_benchmark_dict()))'
29+
30+
json_report_dir="${RUNNER_ARTIFACT_DIR:-${root_dir}}"
31+
json_report_args="--json-report --json-report-file=${json_report_dir}/test-results-libero.json --json-report-indent=2"
32+
33+
python .github/unittest/helpers/coverage_run_parallel.py -m pytest test/libs/test_libero.py \
34+
${json_report_args} \
35+
--instafail -v --durations 200 --capture no \
36+
--timeout 300 --timeout-method=thread \
37+
-k 'not demo_replay_success' \
38+
--error-for-skips
39+
coverage combine -q
40+
coverage xml -i
41+
42+
python .github/unittest/helpers/upload_test_results.py || echo "Warning: Failed to process test results for flaky tracking"

‎.github/unittest/linux_libs/scripts_libero/setup_env.sh‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
#!/usr/bin/env bash
2+
3+
set -euxo pipefail
4+
5+
export DEBIAN_FRONTEND=noninteractive
6+
export TZ="${TZ:-Etc/UTC}"
7+
ln -snf "/usr/share/zoneinfo/${TZ}" /etc/localtime || true
8+
echo "${TZ}" > /etc/timezone || true
9+
10+
apt-get update
11+
apt-get install -y --no-install-recommends tzdata
12+
apt-get upgrade -y
13+
apt-get install -y --no-install-recommends \
14+
cmake \
15+
curl \
16+
ffmpeg \
17+
g++ \
18+
gcc \
19+
git \
20+
libavcodec-dev \
21+
libavdevice-dev \
22+
libavfilter-dev \
23+
libavformat-dev \
24+
libavutil-dev \
25+
libegl1 \
26+
libglew-dev \
27+
libgles2 \
28+
libgl1 \
29+
libglfw3 \
30+
libglvnd0 \
31+
libglx0 \
32+
libosmesa6-dev \
33+
libswresample-dev \
34+
libswscale-dev \
35+
patchelf \
36+
pkg-config \
37+
python3-dev \
38+
unzip \
39+
wget \
40+
xvfb
41+
apt-get upgrade -y libstdc++6
42+
43+
# Ensure EGL discovers the NVIDIA vendor library in the CUDA runner image.
44+
mkdir -p /usr/share/glvnd/egl_vendor.d
45+
cp .github/unittest/linux/scripts/10_nvidia.json /usr/share/glvnd/egl_vendor.d/10_nvidia.json || true
46+
47+
# Avoid error: "fatal: unsafe repository".
48+
git config --global --add safe.directory '*'
49+
50+
root_dir="$(git rev-parse --show-toplevel)"
51+
env_dir="${root_dir}/env"
52+
53+
cd "${root_dir}"
54+
55+
curl -LsSf https://astral.sh/uv/install.sh | sh
56+
export PATH="$HOME/.local/bin:$PATH"
57+
58+
rm -rf "${env_dir}"
59+
uv venv --python "${PYTHON_VERSION}" "${env_dir}"

‎.github/workflows/auto-tag.yml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
# - [Performance], [Perf] → Performance
2626
# - [BC-Breaking], [BCBreaking] → bc breaking
2727
# - [Deprecation], [Deprecated] → Deprecation
28+
# - [Algorithm], [Algorithms] → new algo
2829
# - [Quality] → Quality
2930
# - [Versioning] → versioning
3031
# - [WIP] → WIP
@@ -84,6 +85,7 @@ jobs:
8485
| `[Performance]` or `[Perf]` | Performance | `[Performance] Optimize tensor ops` |
8586
| `[BC-Breaking]` | bc breaking | `[BC-Breaking] Remove deprecated API` |
8687
| `[Deprecation]` | Deprecation | `[Deprecation] Mark old function` |
88+
| `[Algorithm]` or `[Algorithms]` | new algo | `[Algorithm] Add new objective` |
8789
| `[Quality]` | Quality | `[Quality] Fix typos and add codespell` |
8890
| `[Versioning]` | versioning | `[Versioning] Bump release version` |
8991
| `[WIP]` | WIP | `[WIP] Draft implementation` |
@@ -163,6 +165,9 @@ jobs:
163165
Deprecation|Deprecated|deprecation|deprecated)
164166
LABEL="Deprecation"
165167
;;
168+
Algorithm|Algorithms|algorithm|algorithms)
169+
LABEL="new algo"
170+
;;
166171
Quality|quality)
167172
LABEL="Quality"
168173
;;

‎.github/workflows/test-linux-libs.yml‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -367,6 +367,49 @@ jobs:
367367
bash .github/unittest/linux_libs/scripts_jumanji/run_test.sh
368368
bash .github/unittest/linux_libs/scripts_jumanji/post_process.sh
369369
370+
unittests-libero:
371+
strategy:
372+
matrix:
373+
python_version: ["3.10"]
374+
cuda_arch_version: ["12.8"]
375+
fail-fast: false
376+
if: ${{ github.event_name == 'push' || github.event_name == 'workflow_call' || github.event_name == 'workflow_dispatch' || contains(github.event.pull_request.labels.*.name, 'Environments') || contains(github.event.pull_request.labels.*.name, 'Environments/libero') }}
377+
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
378+
with:
379+
repository: pytorch/rl
380+
runner: "linux.g5.4xlarge.nvidia.gpu"
381+
gpu-arch-type: cuda
382+
gpu-arch-version: "12.8"
383+
docker-image: "nvidia/cuda:12.8.0-devel-ubuntu22.04"
384+
timeout: 120
385+
upload-artifact: test-results-libero
386+
script: |
387+
if [[ "${{ github.ref }}" =~ release/* ]]; then
388+
export RELEASE=1
389+
export TORCH_VERSION=stable
390+
else
391+
export RELEASE=0
392+
export TORCH_VERSION=nightly
393+
fi
394+
395+
set -euo pipefail
396+
export PYTHON_VERSION="${{ matrix.python_version }}"
397+
export CU_VERSION="cu128"
398+
export TAR_OPTIONS="--no-same-owner"
399+
export UPLOAD_CHANNEL="nightly"
400+
export TF_CPP_MIN_LOG_LEVEL=0
401+
export BATCHED_PIPE_TIMEOUT=60
402+
export TD_GET_DEFAULTS_TO_NONE=1
403+
404+
nvidia-smi
405+
406+
bash .github/unittest/linux_libs/scripts_libero/setup_env.sh
407+
bash .github/unittest/linux_libs/scripts_libero/install.sh
408+
PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_version.sh "$TORCH_VERSION"
409+
PYTHON=./env/bin/python bash .github/unittest/helpers/assert_torch_tensordict_versions.sh "$TORCH_VERSION"
410+
bash .github/unittest/linux_libs/scripts_libero/run_test.sh
411+
bash .github/unittest/linux_libs/scripts_libero/post_process.sh
412+
370413
unittests-meltingpot:
371414
uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
372415
if: ${{ github.event_name == 'push' || github.event_name == 'workflow_call' || github.event_name == 'workflow_dispatch' || contains(github.event.pull_request.labels.*.name, 'Environments') || contains(github.event.pull_request.labels.*.name, 'Environments/meltingpot') }}

‎benchmarks/test_envs_benchmark.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from tensordict import TensorDict
1111
from torchrl.envs import ParallelEnv, SerialEnv, step_mdp, StepCounter, TransformedEnv
1212
from torchrl.envs.libs.dm_control import DMControlEnv
13+
from torchrl.envs.libs.libero import _has_libero, LiberoEnv
1314
from torchrl.envs.transforms.functional import cat_frames
1415

1516

@@ -41,6 +42,19 @@ def make_parallel_env():
4142
return ((env,), {})
4243

4344

45+
def make_libero_env():
46+
env = LiberoEnv(
47+
"libero_spatial",
48+
task_id=0,
49+
camera_height=256,
50+
camera_width=256,
51+
settle_steps=0,
52+
max_episode_steps=None,
53+
)
54+
env.rollout(3)
55+
return ((env,), {})
56+
57+
4458
def make_nested_td():
4559
return TensorDict(
4660
{
@@ -97,6 +111,13 @@ def test_parallel(benchmark):
97111
benchmark(execute_env, c)
98112

99113

114+
@pytest.mark.skipif(not _has_libero, reason="libero not found")
115+
def test_libero(benchmark):
116+
# raw simulation + render throughput of the LIBERO adapter (steps/s)
117+
(c,), _ = make_libero_env()
118+
benchmark(lambda: c.rollout(100, break_when_any_done=False))
119+
120+
100121
@pytest.mark.parametrize("nested", [True, False])
101122
@pytest.mark.parametrize("keep_other", [True, False])
102123
@pytest.mark.parametrize("exclude_reward", [True, False])

‎docs/source/reference/envs_libraries.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,8 @@ Available wrappers
9494
IsaacLabWrapper
9595
JumanjiEnv
9696
JumanjiWrapper
97+
LiberoEnv
98+
LiberoWrapper
9799
MeltingpotEnv
98100
MeltingpotWrapper
99101
MOGymEnv

0 commit comments

Comments
 (0)