BOLTON-AI/train-robots-with-mujoco

1#!/bin/bash2JUPYTER_TOKEN="${JUPYTER_TOKEN:=huggingface}"3 4NOTEBOOK_DIR="/data/workspaces"5 6# Set up NVIDIA EGL library links at runtime (in case they're mounted differently)7echo "🔗 Setting up NVIDIA EGL library links..."8for nvidia_lib_dir in /usr/local/nvidia/lib64 /usr/local/cuda/lib64 /usr/lib/nvidia; do9    if [ -f "$nvidia_lib_dir/libEGL_nvidia.so.0" ]; then10        echo "Found NVIDIA EGL library at $nvidia_lib_dir/libEGL_nvidia.so.0"11        ln -sf "$nvidia_lib_dir/libEGL_nvidia.so.0" /usr/lib/x86_64-linux-gnu/libEGL_nvidia.so.012        break13    fi14done15 16# perform checks on the GPU configuration17python init_gpu.py18 19# this will download stuff used by Mujoco (the collection of models)20python init_mujoco.py21 22# Copy sample notebooks to persistent storage23echo "📚 Setting up sample workspaces in /data/workspaces..."24 25# Create the workspaces directory if it doesn't exist26mkdir -p /data/workspaces27 28# Create JupyterLab workspace state directory (for UI preferences, layout, etc.)29mkdir -p /data/.jupyter/workspaces30mkdir -p /data/.jupyter/settings31echo "✓ JupyterLab workspace state directory: /data/.jupyter/workspaces"32echo "✓ JupyterLab settings directory: /data/.jupyter/settings"33 34# Loop through each sample directory35for sample_dir in $HOME/app/samples/*/; do36    # Get the sample name (e.g., "locomotion" from "samples/locomotion/")37    sample_name=$(basename "$sample_dir")38 39    # Create the workspace directory if it doesn't exist40    workspace_dir="/data/workspaces/$sample_name"41    if [ ! -d "$workspace_dir" ]; then42        echo "  Creating workspace: $workspace_dir"43        mkdir -p "$workspace_dir"44    else45        echo "  Workspace already exists: $workspace_dir"46    fi47 48    # Copy the .ipynb file only if it doesn't already exist (to preserve user changes)49    ipynb_file="$sample_dir$sample_name.ipynb"50    dest_ipynb="$workspace_dir/$sample_name.ipynb"51 52    if [ -f "$ipynb_file" ]; then53        if [ ! -f "$dest_ipynb" ]; then54            echo "  Copying: $sample_name.ipynb → $workspace_dir/"55            cp "$ipynb_file" "$dest_ipynb"56        else57            echo "  Preserving existing: $dest_ipynb (user may have made changes)"58        fi59    fi60 61    # Copy any other files from the sample directory (excluding .ipynb files)62    # This allows samples to include datasets, scripts, etc.63    for file in "$sample_dir"*; do64        filename=$(basename "$file")65        # Skip if it's the .ipynb file (already handled) or if it's a directory66        if [ "$filename" != "$sample_name.ipynb" ] && [ -f "$file" ]; then67            dest_file="$workspace_dir/$filename"68            if [ ! -f "$dest_file" ]; then69                echo "  Copying additional file: $filename → $workspace_dir/"70                cp "$file" "$dest_file"71            fi72        fi73    done74done75 76echo "✅ Sample workspaces ready!"77echo ""78 79# Initialize OpenTrack (download datasets, create symlinks)80if [ -f "$HOME/app/samples/opentrack/init_opentrack.sh" ]; then81    bash "$HOME/app/samples/opentrack/init_opentrack.sh"82fi83 84jupyter labextension disable "@jupyterlab/apputils-extension:announcements"85 86jupyter-lab \87    --ip 0.0.0.0 \88    --port 7860 \89    --no-browser \90    --allow-root \91    --ServerApp.token="$JUPYTER_TOKEN" \92    --ServerApp.tornado_settings="{'headers': {'Content-Security-Policy': 'frame-ancestors *'}}" \93    --ServerApp.cookie_options="{'SameSite': 'None', 'Secure': True}" \94    --ServerApp.disable_check_xsrf=True \95    --LabApp.news_url=None \96    --LabApp.check_for_updates_class="jupyterlab.NeverCheckForUpdate" \97    --notebook-dir=$NOTEBOOK_DIR \98    --ServerApp.default_url="/data/workspaces/opentrack/opentrack_quickstart.ipynb"99