setup.sh

#!/bin/bash
# One-time (idempotent) setup for the emdb-simulator workspace.
# Safe to re-run: every step checks whether it's already done before acting.
#
# Usage:
#   ./setup.sh            # full setup
#   ./setup.sh --docs      # also install docs/ build dependencies
#
# Override the venv location (default: <repo>/.venv) with:
#   VENV_DIR=/path/to/venv ./setup.sh
set -euo pipefail

# ---------------------------------------------------------------------------
# 0. Locate ourselves. No hardcoded machine paths anywhere in this script.
# ---------------------------------------------------------------------------
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV_DIR="${VENV_DIR:-$REPO_ROOT/.venv}"
WITH_DOCS=0
for arg in "$@"; do
    case "$arg" in
        --docs) WITH_DOCS=1 ;;
        *) echo "Unknown argument: $arg" >&2; exit 1 ;;
    esac
done

log() { echo -e "\n==> $*"; }

cd "$REPO_ROOT"

# ---------------------------------------------------------------------------
# 1. ROS 2: auto-detect the installed distro instead of hardcoding one.
# ---------------------------------------------------------------------------
log "Locating ROS 2 installation..."
if [ -z "${ROS_DISTRO:-}" ]; then
    for d in /opt/ros/*/; do
        distro="$(basename "$d")"
        if [ -f "$d/setup.bash" ]; then
            ROS_DISTRO="$distro"
            break
        fi
    done
fi
if [ -z "${ROS_DISTRO:-}" ] || [ ! -f "/opt/ros/$ROS_DISTRO/setup.bash" ]; then
    echo "No ROS 2 installation found under /opt/ros. Install ROS 2 first: https://docs.ros.org/en/rolling/Installation.html" >&2
    exit 1
fi
echo "Using ROS 2 distro: $ROS_DISTRO"
# ROS 2's setup.bash references unset variables (e.g. AMENT_TRACE_SETUP_FILES),
# which is incompatible with `set -u`; disable it just for the source.
set +u
# shellcheck disable=SC1090
source "/opt/ros/$ROS_DISTRO/setup.bash"
set -u

# ---------------------------------------------------------------------------
# 2. Git submodules (misc/robosuite, misc/robocasa, misc/robosuite_models).
#    Only initialize paths that are actually empty, so this doesn't clobber
#    a manually-checked-out working copy that isn't wired up as a proper
#    submodule worktree yet.
# ---------------------------------------------------------------------------
log "Checking git submodules..."
if [ -f .gitmodules ]; then
    while read -r path; do
        [ -z "$path" ] && continue
        if [ -d "$path" ] && [ -n "$(ls -A "$path" 2>/dev/null)" ]; then
            echo "  $path already populated, skipping"
        else
            echo "  initializing $path"
            git submodule update --init --recursive -- "$path"
        fi
    done < <(git config -f .gitmodules --get-regexp path | awk '{print $2}')
fi

# ---------------------------------------------------------------------------
# 3. Python virtualenv.
# ---------------------------------------------------------------------------
log "Setting up Python virtualenv at $VENV_DIR ..."
if [ ! -d "$VENV_DIR" ]; then
    # --system-site-packages: the ROS 2 build toolchain (rosidl's lark parser,
    # colcon, etc.) is installed system-wide via apt, not pip. Without this,
    # colcon build fails deep inside rosidl with "No module named 'lark'".
    # Packages pip-installs into the venv (numpy, robosuite, ...) still take
    # priority over the system copies.
    python3 -m venv --system-site-packages "$VENV_DIR"
fi
# shellcheck disable=SC1091
source "$VENV_DIR/bin/activate"
python -m pip install --upgrade pip wheel setuptools

log "Installing torch..."
# robocasa depends on lerobot==0.3.3, which pulls in torch<2.8.0,>=2.2.1;
# pin the exact version explicitly (numeric-version parity between the CPU
# and GPU cases below) rather than leaving it to whatever pip's resolver
# picks that day.
#
# Machines without a working NVIDIA GPU+driver get the CPU-only wheel:
# otherwise pip pulls torch's default CUDA build -- multiple GB of unneeded
# nvidia-*-cu12 wheels with nothing to run them on, and on machines whose
# installed driver doesn't match that build's expected CUDA runtime, an
# import-time crash (missing libcudartX.so). Same fix as
# docker/Dockerfile.cpu and .github/workflows/docs.yml.
# Machines that DO have a working GPU get the CUDA build (matches
# docker/Dockerfile.gpu) so RoboCasa/train_sb3 can actually use it.
if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi >/dev/null 2>&1; then
    echo "  NVIDIA GPU detected -- installing CUDA-enabled torch."
    pip install torch==2.7.1 torchvision==0.22.1
else
    echo "  No usable NVIDIA GPU detected -- installing CPU-only torch."
    echo "  (train_sb3's PPO policy is a tiny MLP; the sim/ROS round-trip is"
    echo "  the actual bottleneck, so this is fine even for training. If"
    echo "  this detection is wrong for your machine -- e.g. nvidia-smi"
    echo "  fails because the driver needs a reboot after an update -- pip"
    echo "  install the CUDA build manually once that's sorted out.)"
    pip install torch==2.7.1+cpu torchvision==0.22.1+cpu \
        --index-url https://download.pytorch.org/whl/cpu
fi

log "Installing robosuite (editable)..."
pip install -e "$REPO_ROOT/misc/robosuite"

log "Installing robocasa (editable)..."
pip install -e "$REPO_ROOT/misc/robocasa"

if [ -n "$(ls -A "$REPO_ROOT/misc/robosuite_models" 2>/dev/null)" ]; then
    log "Installing robosuite_models (editable, optional)..."
    pip install -e "$REPO_ROOT/misc/robosuite_models" || \
        echo "  robosuite_models install failed (optional dependency) -- continuing"
else
    echo "  misc/robosuite_models is empty, skipping (optional dependency)"
fi

log "Installing gymnasium + stable-baselines3 (emdb_policy's train_sb3)..."
pip install gymnasium stable-baselines3

if [ "$WITH_DOCS" -eq 1 ]; then
    log "Installing docs/ build dependencies..."
    pip install -r "$REPO_ROOT/docs/requirements.txt"
fi

log "Checking RoboCasa kitchen assets (textures/fixtures/objects)..."
# robocasa prints its own stdout noise on import (e.g. a "mimicgen not
# installed" notice), so the result is tagged with a unique marker and
# extracted with grep/sed rather than trusting stdout to only contain it.
MISSING_ASSET_TYPES="$(python3 - <<'PYEOF' | grep '^MISSING_ASSET_TYPES:' | sed 's/^MISSING_ASSET_TYPES://'
import os
from robocasa.scripts.download_kitchen_assets import DOWNLOAD_ASSET_REGISTRY

missing = [
    name for name, cfg in DOWNLOAD_ASSET_REGISTRY.items()
    if not (os.path.isdir(cfg["folder"]) and os.listdir(cfg["folder"]))
]
print(f"MISSING_ASSET_TYPES:{' '.join(missing)}")
PYEOF
)"
if [ -n "$MISSING_ASSET_TYPES" ]; then
    echo "  Missing/empty: $MISSING_ASSET_TYPES (downloading, ~10GB total -- only the missing types)"
    # download_kitchen_assets asks for a one-time y/n confirmation before
    # fetching; auto-confirm it instead of hanging non-interactively.
    yes | python -m robocasa.scripts.download_kitchen_assets --type $MISSING_ASSET_TYPES
else
    echo "  All kitchen asset types already present, skipping download."
fi

# colcon must be importable via the venv's python (python -m colcon), so that
# packages built by it get console-script shebangs pointing at the venv --
# otherwise built nodes can't import venv-only deps like robosuite.
log "Installing colcon into the venv..."
pip install colcon-common-extensions
# colcon-common-extensions pulls in the latest EmPy (4.x), but ROS 2 Humble's
# rosidl_adapter only works with the old em 3.3.4 API; pin it back down.
pip install "empy==3.3.4"

# ---------------------------------------------------------------------------
# 4. ROS dependencies via rosdep.
# ---------------------------------------------------------------------------
log "Resolving ROS package dependencies via rosdep..."
if [ ! -f /etc/ros/rosdep/sources.list.d/20-default.list ]; then
    echo "  rosdep has not been initialized on this machine; running 'sudo rosdep init'"
    sudo rosdep init
fi
rosdep update
rosdep install --from-paths "$REPO_ROOT/ros_packages/src" --ignore-src -r -y

# ---------------------------------------------------------------------------
# 5. Build the ROS 2 workspace.
# ---------------------------------------------------------------------------
log "Building ros_packages workspace..."
(
    cd "$REPO_ROOT/ros_packages"
    python -m colcon build --symlink-install
)

log "Setup complete."
cat <<EOF

Repo root : $REPO_ROOT
Venv      : $VENV_DIR
ROS distro: $ROS_DISTRO

For every new terminal, run:
    source $REPO_ROOT/env.sh
EOF