install.sh

#!/usr/bin/env bash
#
# Create the conda environment for UDP and build rbdl with its Python bindings,
# compiled against this environment's interpreter.
#
#   ./install.sh                 # create env (default name: UIP) + build rbdl
#   ./install.sh --name myenv    # use a different env name
#   ./install.sh --no-rbdl       # only create the python env
#   ./install.sh --rbdl-only     # only (re)build rbdl into an existing env
#
# Env overrides: ENV_NAME, PYTHON_VERSION, TORCH_VERSION, TORCHVISION_VERSION,
#                TORCH_INDEX_URL, RBDL_SRC, RECREATE=1 (remove existing env first).
#
# NOTE: the repo's other scripts (eval_scripts/, train_scripts/, data_preprocessing/)
# default to UDP_ENV=UDP, so the env is named "UDP" by default. If you create it under
# another name, run those scripts with UDP_ENV=<name>.
#
set -eo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$REPO_ROOT"

ENV_NAME="${ENV_NAME:-UDP}"
PYTHON_VERSION="${PYTHON_VERSION:-3.8.10}"
TORCH_VERSION="${TORCH_VERSION:-2.3.1}"
TORCHVISION_VERSION="${TORCHVISION_VERSION:-0.18.1}"
TORCH_INDEX_URL="${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu121}"
RBDL_SRC="${RBDL_SRC:-$(dirname "$REPO_ROOT")/rbdl}"   # default: sibling of the repo (…/repos/rbdl)
DO_ENV=1; DO_RBDL=1

usage() { sed -n '3,16p' "$0" | sed 's/^# \{0,1\}//'; }
while [[ $# -gt 0 ]]; do
    case "$1" in
        --name)      ENV_NAME="$2"; shift 2 ;;
        --name=*)    ENV_NAME="${1#*=}"; shift ;;
        --no-rbdl)   DO_RBDL=0; shift ;;
        --rbdl-only) DO_ENV=0; shift ;;
        -h|--help)   usage; exit 0 ;;
        *) echo "Unknown argument: $1"; usage; exit 1 ;;
    esac
done

command -v conda >/dev/null 2>&1 || { echo "ERROR: conda not found on PATH." >&2; exit 1; }
eval "$(conda shell.bash hook)"

# ---------------------------------------------------------------- python env ---
if [ "$DO_ENV" = 1 ]; then
    if conda env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then
        if [ "${RECREATE:-0}" = 1 ]; then
            echo ">>> Removing existing env '$ENV_NAME' (RECREATE=1)"
            conda env remove -y -n "$ENV_NAME"
        else
            echo "ERROR: conda env '$ENV_NAME' already exists. Use RECREATE=1 to replace it," >&2
            echo "       pass --name <other>, or --rbdl-only to skip env creation." >&2
            exit 1
        fi
    fi

    echo ">>> Creating conda env '$ENV_NAME' (python $PYTHON_VERSION)"
    conda create -y -n "$ENV_NAME" "python=$PYTHON_VERSION"
    conda activate "$ENV_NAME"
    python -m pip install --upgrade pip wheel setuptools

    echo ">>> Installing PyTorch $TORCH_VERSION / torchvision $TORCHVISION_VERSION ($TORCH_INDEX_URL)"
    pip install "torch==$TORCH_VERSION" "torchvision==$TORCHVISION_VERSION" --index-url "$TORCH_INDEX_URL"

    echo ">>> Installing numpy + cython first (chumpy builds against them)"
    pip install numpy==1.23.1 cython==0.29.30

    echo ">>> Installing chumpy (needs numpy present; SMPL .pkl loading depends on it)"
    pip install --no-build-isolation chumpy==0.70

    echo ">>> Installing the rest of requirements.txt"
    pip install -r requirements.txt

    echo ">>> Import sanity check"
    python - <<'PY'
import torch, numpy, scipy, matplotlib, aitviewer, smplx, chumpy, wandb
import cv2, plotly, vctoolkit, pybullet, qpsolvers, einops, easydict, prettytable
print("torch", torch.__version__, "| cuda", torch.version.cuda, "| cuda available:", torch.cuda.is_available())
print("aitviewer", aitviewer.__version__, "| numpy", numpy.__version__)
print("core imports OK")
PY
else
    conda activate "$ENV_NAME"
fi

# ----------------------------------------------------- activation env vars ---
# Persist the env vars the repo needs so they are set automatically on every
# `conda activate $ENV_NAME` (no need to export PYTHONPATH / SMPLX_MODEL_PATH by
# hand). Written as an activate.d hook so PYTHONPATH is appended, not clobbered.
ACTIVATE_D="$CONDA_PREFIX/etc/conda/activate.d"
mkdir -p "$ACTIVATE_D"
cat > "$ACTIVATE_D/udp_env.sh" <<EOF
# Auto-generated by install.sh — UDP repo environment variables.
export PYTHONPATH="$REPO_ROOT:\$PYTHONPATH"
export SMPLX_MODEL_PATH="$REPO_ROOT/data/smpl_m_lbs_10_207_0_v1.0.0.pkl"
EOF
echo ">>> Wrote activation hook: $ACTIVATE_D/udp_env.sh"
echo "    (sets PYTHONPATH + SMPLX_MODEL_PATH on every 'conda activate $ENV_NAME')"

# -------------------------------------------------------------------- rbdl ---
if [ "$DO_RBDL" = 1 ]; then
    echo ">>> Building rbdl with Python bindings into env '$ENV_NAME'"

    # Build prerequisites: a recent cmake + Eigen3 (conda, no sudo), Cython + numpy (in the env).
    conda install -y -n "$ENV_NAME" -c conda-forge cmake eigen
    pip install --quiet cython==0.29.30 numpy==1.23.1
    if ! command -v c++ >/dev/null 2>&1 && ! command -v g++ >/dev/null 2>&1; then
        echo "ERROR: no C++ compiler found. Install one (e.g. 'sudo apt install build-essential')" >&2
        echo "       or 'conda install -n $ENV_NAME -c conda-forge cxx-compiler', then re-run --rbdl-only." >&2
        exit 1
    fi

    if [ -d "$RBDL_SRC/.git" ]; then
        echo ">>> rbdl already cloned at $RBDL_SRC"
    else
        echo ">>> Cloning rbdl (recursive) into $RBDL_SRC"
        git clone --recursive https://github.com/rbdl/rbdl "$RBDL_SRC"
    fi

    BUILD_DIR="$RBDL_SRC/rbdl-build"
    PYEXE="$(python -c 'import sys; print(sys.executable)')"
    EIGEN_INC="$CONDA_PREFIX/include/eigen3"
    mkdir -p "$BUILD_DIR"
    echo ">>> Configuring rbdl against $PYEXE"
    cmake -S "$RBDL_SRC" -B "$BUILD_DIR" \
        -D CMAKE_BUILD_TYPE=Release \
        -D RBDL_BUILD_PYTHON_WRAPPER=ON \
        -D RBDL_BUILD_ADDON_URDFREADER=ON \
        -D PYTHON_EXECUTABLE="$PYEXE" \
        -D Python_EXECUTABLE="$PYEXE" \
        -D Python3_EXECUTABLE="$PYEXE" \
        -D CMAKE_PREFIX_PATH="$CONDA_PREFIX" \
        -D EIGEN3_INCLUDE_DIR="$EIGEN_INC"
    cmake --build "$BUILD_DIR" --target all -j "$(nproc)"

    RBDL_PY="$BUILD_DIR/python"
    echo ">>> Verifying rbdl python bindings"
    if PYTHONPATH="$RBDL_PY" python -c "import rbdl; print('rbdl OK:', rbdl.__file__)"; then
        echo
        echo "rbdl built successfully. To use the physics optimizer, export (matches"
        echo "eval_scripts/_common.sh's RBDL_PYTHON_PATH):"
        echo "    export RBDL_PYTHON_PATH=$RBDL_PY"
    else
        echo "WARNING: rbdl built but failed to import — check the build log above." >&2
    fi
fi

echo
echo "Done. Activate with:  conda activate $ENV_NAME"