[Feature] Add support for loading datasets from local Minari cache (#3068)

Co-authored-by: vmoens <vincentmoens@gmail.com>

Author: Ibinarriaga8

.github/unittest/linux_libs/scripts_minari/environment.yml

@@ -20,4 +20,11 @@ dependencies:
     - scipy
     - hydra-core
     - minari[gcs,hdf5,hf]
-    - gymnasium<1.0.0
+    - gymnasium>=1.2.0
+    - ale-py
+    - gymnasium-robotics
+    - minari[create]
+    - jax
+    - mujoco               
+    - mujoco-py<2.2,>=2.1
+    - minigrid

.github/unittest/linux_libs/scripts_minari/install.sh

@@ -7,8 +7,7 @@ unset PYTORCH_VERSION
 
 set -e
 
-eval "$(./conda/bin/conda shell.bash hook)"
-conda activate ./env
+# Note: This script is sourced by run_all.sh, so the environment is already active
 
 if [ "${CU_VERSION:-}" == cpu ] ; then
     version="cpu"
@@ -22,22 +21,21 @@ else
     version="$(python -c "print('.'.join(\"${CUDA_VERSION}\".split('.')[:2]))")"
 fi
 
-
 # submodules
 git submodule sync && git submodule update --init --recursive
 
 printf "Installing PyTorch with cu128"
 if [[ "$TORCH_VERSION" == "nightly" ]]; then
   if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U
+      uv pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U
   else
-      pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128 -U
+      uv pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128 -U
   fi
 elif [[ "$TORCH_VERSION" == "stable" ]]; then
     if [ "${CU_VERSION:-}" == cpu ] ; then
-      pip3 install torch --index-url https://download.pytorch.org/whl/cpu
+      uv pip install torch --index-url https://download.pytorch.org/whl/cpu
   else
-      pip3 install torch --index-url https://download.pytorch.org/whl/cu128
+      uv pip install torch --index-url https://download.pytorch.org/whl/cu128
   fi
 else
   printf "Failed to install pytorch"
@@ -46,9 +44,9 @@ fi
 
 # install tensordict
 if [[ "$RELEASE" == 0 ]]; then
-  pip3 install git+https://github.com/pytorch/tensordict.git
+  uv pip install git+https://github.com/pytorch/tensordict.git
 else
-  pip3 install tensordict
+  uv pip install tensordict
 fi
 
 # smoke test

.github/unittest/linux_libs/scripts_minari/post_process.sh

@@ -2,5 +2,4 @@
 
 set -e
 
-eval "$(./conda/bin/conda shell.bash hook)"
-conda activate ./env
+# Note: This script is sourced by run_all.sh, so the environment is already active

.github/unittest/linux_libs/scripts_minari/requirements.txt

@@ -0,0 +1,20 @@
+hypothesis
+future
+cloudpickle
+pytest
+pytest-cov
+pytest-mock
+pytest-instafail
+pytest-rerunfailures
+pytest-error-for-skips
+pytest-asyncio
+expecttest
+pybind11[global]
+pyyaml
+scipy
+hydra-core
+minari[gcs,hdf5,hf,create]
+gymnasium>=1.2.0
+ale-py
+gymnasium-robotics
+mujoco 

.github/unittest/linux_libs/scripts_minari/run_all.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+
+set -euxo pipefail
+
+# Run all minari test scripts in sequence, sourcing each one to maintain environment state
+this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+
+echo "Running minari tests with uv-based setup..."
+
+# Source each script in sequence to maintain environment state
+source "${this_dir}/setup_env.sh"
+source "${this_dir}/install.sh"
+source "${this_dir}/run_test.sh"
+source "${this_dir}/post_process.sh"
+
+echo "Minari tests completed successfully!" 

.github/unittest/linux_libs/scripts_minari/run_test.sh

@@ -2,8 +2,7 @@
 
 set -e
 
-eval "$(./conda/bin/conda shell.bash hook)"
-conda activate ./env
+# Note: This script is sourced by run_all.sh, so the environment is already active
 
 apt-get update && apt-get remove swig -y && apt-get install -y git gcc patchelf libosmesa6-dev libgl1-mesa-glx libglfw3 swig3.0 cmake
 ln -s /usr/bin/swig3.0 /usr/bin/swig
@@ -18,8 +17,6 @@ root_dir="$(git rev-parse --show-toplevel)"
 env_dir="${root_dir}/env"
 lib_dir="${env_dir}/lib"
 
-conda deactivate && conda activate ./env
-
 # this workflow only tests the libs
 python -c "import minari"
 

.github/unittest/linux_libs/scripts_minari/setup_env.sh

@@ -33,37 +33,40 @@ apt-get upgrade -y libstdc++6
 
 this_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 root_dir="$(git rev-parse --show-toplevel)"
-conda_dir="${root_dir}/conda"
 env_dir="${root_dir}/env"
 
 cd "${root_dir}"
 
-case "$(uname -s)" in
-    Darwin*) os=MacOSX;;
-    *) os=Linux
-esac
-
-# 1. Install conda at ./conda
-if [ ! -d "${conda_dir}" ]; then
-    printf "* Installing conda\n"
-    wget -O miniconda.sh "http://repo.continuum.io/miniconda/Miniconda3-latest-${os}-x86_64.sh"
-    bash ./miniconda.sh -b -f -p "${conda_dir}"
+# Install uv if not already installed
+if ! command -v uv &> /dev/null; then
+    printf "* Installing uv\n"
+    # Try different Python commands
+    if command -v python3 &> /dev/null; then
+        python3 -m pip install uv
+    elif command -v python &> /dev/null; then
+        python -m pip install uv
+    else
+        # Fallback to curl installation
+        curl -LsSf https://astral.sh/uv/install.sh | sh
+        export PATH="$HOME/.local/bin:$PATH"
+    fi
 fi
-eval "$(${conda_dir}/bin/conda shell.bash hook)"
 
-# 2. Create test environment at ./env
+# Create virtual environment using uv
 printf "python: ${PYTHON_VERSION}\n"
 if [ ! -d "${env_dir}" ]; then
-    printf "* Creating a test environment\n"
-    conda create --prefix "${env_dir}" -y python="$PYTHON_VERSION"
+    printf "* Creating a test environment with uv\n"
+    uv venv "${env_dir}" --python "${PYTHON_VERSION}"
 fi
-conda activate "${env_dir}"
 
-# 3. Install Conda dependencies
-printf "* Installing dependencies (except PyTorch)\n"
-echo "  - python=${PYTHON_VERSION}" >> "${this_dir}/environment.yml"
-cat "${this_dir}/environment.yml"
+# Activate the virtual environment
+source "${env_dir}/bin/activate"
 
-pip3 install pip --upgrade
+# Upgrade pip
+uv pip install --upgrade pip
 
-conda env update --file "${this_dir}/environment.yml" --prune
+# Install dependencies from requirements.txt (we'll create this)
+printf "* Installing dependencies (except PyTorch)\n"
+if [ -f "${this_dir}/requirements.txt" ]; then
+    uv pip install -r "${this_dir}/requirements.txt"
+fi

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

@@ -466,10 +466,7 @@ jobs:
         export BATCHED_PIPE_TIMEOUT=60
         export TD_GET_DEFAULTS_TO_NONE=1
 
-        bash .github/unittest/linux_libs/scripts_minari/setup_env.sh
-        bash .github/unittest/linux_libs/scripts_minari/install.sh
-        bash .github/unittest/linux_libs/scripts_minari/run_test.sh
-        bash .github/unittest/linux_libs/scripts_minari/post_process.sh
+        bash .github/unittest/linux_libs/scripts_minari/run_all.sh
 
   unittests-openx:
     strategy:

packaging/build_wheels.sh

@@ -7,7 +7,7 @@ script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
 export BUILD_TYPE=wheel
 setup_env
 setup_wheel_python
-pip_install numpy pyyaml future ninja
+pip_install numpy pyyaml future ninja "pybind11[global]"
 pip_install --upgrade setuptools
 setup_pip_pytorch_version
 python setup.py clean

test/_utils_internal.py

@@ -13,6 +13,7 @@
 import unittest
 import warnings
 from functools import wraps
+from typing import Callable
 
 import pytest
 import torch
@@ -214,7 +215,12 @@ def generate_seeds(seed, repeat):
 
 
 # Decorator to retry upon certain Exceptions.
-def retry(ExceptionToCheck, tries=3, delay=3, skip_after_retries=False):
+def retry(
+    ExceptionToCheck: type[Exception],
+    tries: int = 3,
+    delay: int = 3,
+    skip_after_retries: bool = False,
+) -> Callable[[Callable], Callable]:
     def deco_retry(f):
         @wraps(f)
         def f_retry(*args, **kwargs):