# SkyPilot task for launching mjlab training on Lambda Cloud.
#
# Uses the pre-built Docker image from GHCR.
#
# Usage:
#   sky launch scripts/cloud/train-docker.yaml \
#     --env TASK=Mjlab-Velocity-Flat-Unitree-G1

name: mjlab-train

resources:
  cloud: lambda
  accelerators: A100:1
  autostop:
    idle_minutes: 5
    down: true  # Terminates the instance when idle (stops billing).

workdir: .

file_mounts:
  ~/.netrc: ~/.netrc

envs:
  TASK: Mjlab-Velocity-Flat-Unitree-G1
  NUM_ENVS: "4096"
  MAX_ITERATIONS: "6000"
  MUJOCO_GL: egl

setup: |
  # Configure NVIDIA runtime for Docker if not already set up.
  if ! sudo docker info 2>/dev/null | grep -q "nvidia"; then
    sudo nvidia-ctk runtime configure --runtime=docker
    sudo systemctl restart docker
    sleep 3  # Wait for the daemon to be ready before pulling.
  fi

  sudo docker pull ghcr.io/mujocolab/mjlab:latest

run: |
  sudo docker run --rm --runtime=nvidia --gpus all \
    -v "$HOME/.netrc:/root/.netrc:ro" \
    -e MUJOCO_GL=egl \
    ghcr.io/mujocolab/mjlab:latest \
    uv run --no-dev train "$TASK" \
      --env.scene.num-envs "$NUM_ENVS" \
      --agent.max-iterations "$MAX_ITERATIONS" \
      --gpu-ids all