# 
# This is the Dockerfile for the initial setup of the 
# SO-101 robots. It's meant to showcase a clean environment
# to be used along with the ARC documentation as a sort 
# of "Hello World".
#
# Additionally, this container can be used as a base image for 
# further research, as is often done within ARC. But this file 
# is just to get up and running. You should not edit this, but 
# rather extend it within your experiments 
#
#
# BUILDING
# ===========================================================
#
#       docker build -t so101-base .
#
#
# CREATING THE CONTAINER
# ===========================================================
# 
# Use the below flags:
#
#       --gpus all: Grants the container access to the 
#                   GPUs on your system.
#
#       --privileged and -v /dev:/dev: Crucial for capturing 
#                   data streams from cameras (/dev/video*) and 
#                   communicating with the SO-101 motor controllers 
#                   (/dev/ttyUSB* or /dev/ttyACM*).
#
#       --shm-size=128g: Prevents "bus error" crashes during 
#                   training by providing enough shared memory 
#                   for large robotic datasets. 
#
#       - v /run/user/$(id -u)/pulse:
#                   The last three with pulse, they are to allowing access
#                   for audio from the container to be played on the host
#
#
#       docker run -it \
#          --name so101 \
#          --env="DISPLAY=$DISPLAY" \
#          --env="XAUTHORITY=$XAUTHORITY" \
#          --volume="$XAUTHORITY:$XAUTHORITY" \
#          --volume="/tmp/.X11-unix:/tmp/.X11-unix:rw" \
#          --gpus all \
#          --privileged \
#          --network host \
#          --shm-size=128g \
#          -v ~/code/robotics/so-101/resources/LeRobot:/LeRobot \
#          -v /dev:/dev \
#          -v /run/user/$(id -u)/pulse:/run/user/$(id -u)/pulse \
#          -e PULSE_SERVER=unix:/run/user/$(id -u)/pulse/native \
#          -v ~/.config/pulse/cookie:/root/.config/pulse/cookie \
#          so101-base bash
#        
#
# VERIFICATION
# ===========================================================
# 
# Once inside the container, run these commands to confirm 
# everything is connected:
#
#   NVIDIA Error:
#       If you see this, run the below python code to verify all is well: 
#    
#           "WARNING: Detected NVIDIA GeForce RTX 5060 Ti GPU, which is 
#           not yet supported in this version of the container
#           ERROR: No supported GPU(s) detected to run this container"
#
#       # This will confirm if you need to worry or not...
#       python3 -c "import torch; print(f'Device: {torch.cuda.get_device_name(0)}'); print(f'CUDA working: {torch.cuda.is_available()}'); x = torch.ones(1).cuda(); print('Tensor on GPU: Success')"
#
#
#   GPU Availablility: 
#
#       nvidia-smi 
#
#
#    Cameras: ensure cameras are detected
#
#           v4l2-ctl --list-devices 
#
#    Robot: to find your SO-101 ports. 
#
#           ls /dev/ttyACM* 
#               - or - 
#           ls /dev/ttyUSB*  
#
#
# DOCKER DEFINITION
# ===========================================================
# Use the newer NVIDIA image that includes CUDA support and Python 3.12
FROM nvcr.io/nvidia/pytorch:24.12-py3

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

# Add a User with Sudo rights:
# that matches your host ID but has "passwordless sudo" access:
ARG USER_ID=1000
ARG GROUP_ID=1000

# 1. Update apt and install sudo
# 2. Check if UID/GID 1000 exists and delete them to avoid conflicts
# 3. Create so101-user and set permissions
RUN apt-get update && apt-get install -y sudo && \
    if getent passwd ${USER_ID} ; then userdel -f $(getent passwd ${USER_ID} | cut -d: -f1); fi && \
    if getent group ${GROUP_ID} ; then groupdel $(getent group ${GROUP_ID} | cut -d: -f1); fi && \
    groupadd -g ${GROUP_ID} so101-user && \
    useradd -m -u ${USER_ID} -g so101-user -s /bin/bash so101-user && \
    usermod -aG dialout,tty so101-user && \
    echo "so101-user ALL=(ALL) NOPASSWD:ALL" >> /etc/sudoers && \
    rm -rf /var/lib/apt/lists/*

# Install system dependencies for the so-101, cameras, and robotics
RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    cmake \
    build-essential \
    libgl1 \
    libglx-mesa0 \
    libglib2.0-0 \
    ffmpeg \
    v4l-utils \
    usbutils \
    speech-dispatcher \
    speech-dispatcher-espeak-ng \
    espeak-ng \
    pulseaudio-utils \
    && rm -rf /var/lib/apt/lists/*


# Set the working directory
WORKDIR /workspace

# The NVIDIA PyTorch base image pre-creates a /workspace directory with 
# some default NVIDIA readme files and examples inside it.
# We need to remove those, then clone the LeRobot repo. if we don't remove them 
# first, this fails
#
# This base image preinstalls LeRobot as it is the simplest to get up and running quickly,
# validate your setup, etc. It's unlikely you will use LeRobot in production environments.
RUN rm -rf ./* && \
    git clone https://github.com/huggingface/lerobot.git . && \
    git checkout 4eecbad32b0380c64bd5574bbf74e56e5255bdea && \
    pip install --no-cache-dir --upgrade pip && \
    # 1. Pre-install the core conflicting libraries with NumPy 1.x compatible versions
    pip install --no-cache-dir \
        "numpy<2.0.0" \
        "opencv-python-headless<4.10" \
        "pyrealsense2" && \
    # 2. Install LeRobot. 
    # IMPORTANT: We use -e to install the local repo, 
    # but we must ensure it doesn't pull in a newer NumPy.
    pip install --no-cache-dir \
        transformers \
        accelerate \
        timm \
        einops \
        num2words \
        -e ".[feetech, dynamixel, aloha, intelrealsense, async]" && \
    # 3. Double-check that NumPy didn't get snuck back in by another dependency
    pip install --no-cache-dir "numpy<2.0.0"


# Set the working directory
WORKDIR /LeRobot

# Set the LeRobot home directory to your mounted volume path
ENV HF_LEROBOT_HOME=/LeRobot

USER so101-user


# Entry point
CMD ["/bin/bash"]