#!/usr/bin/env bash
# GRPO on the GridSim v1 bundle through AgentEnv's agent function, as the run the Miles page plots ran it on four L40S
# GPUs in radixark/miles:dev-202611020756: the Qwen3-4B SFT warm start as a full fine-tune on FSDP, colocated with
# SGLang, lr 1e-6 constant, 32 rollouts a step (4 days x 8) at temperature 1.1, thinking off, at most 4 model turns a
# rollout, and evals on the 200 eval days at steps 0, 10 and 20. Run it in Miles' container from a folder holding
# qwen3-4b-gridsim-sft and the rows python -m agent_env.training.miles wrote to data/, with agentenv-framework[miles]
# installed; with a server that asks for a token, export the variable the rows' --token-var names first.
# Miles' launcher refuses GPUs it has no verified profile for, as it did the L40S, so this starts Ray and submits
# train.py as the launcher does.
set -euo pipefail
ray start --head --node-ip-address 127.0.0.1 --num-gpus 4 --disable-usage-stats
# The job runs where Ray puts it, so the script starts in this folder.
{ echo "cd $(printf %q "$PWD")"; cat <<'EOF'
exec python3 /root/miles/train.py \
  --hf-checkpoint qwen3-4b-gridsim-sft \
  --prompt-data data/train.jsonl \
  --input-key messages \
  --metadata-key metadata \
  --rollout-shuffle \
  --apply-chat-template-kwargs '{"enable_thinking": false}' \
  --num-rollout 20 \
  --rollout-batch-size 4 \
  --n-samples-per-prompt 8 \
  --global-batch-size 32 \
  --rollout-temperature 1.1 \
  --rollout-max-response-len 1024 \
  --max-seq-len 8192 \
  --eval-prompt-data eval data/eval.jsonl \
  --eval-interval 10 \
  --n-samples-per-eval-prompt 1 \
  --eval-temperature 1.1 \
  --eval-max-response-len 1024 \
  --advantage-estimator grpo \
  --kl-loss-coef 0.0 \
  --kl-coef 0.0 \
  --entropy-coef 0.0 \
  --eps-clip 0.2 \
  --eps-clip-high 0.28 \
  --optimizer adam \
  --lr 1e-6 \
  --lr-decay-style constant \
  --weight-decay 0.0 \
  --adam-beta1 0.9 \
  --adam-beta2 0.98 \
  --train-backend fsdp \
  --gradient-checkpointing \
  --update-weight-buffer-size 536870912 \
  --actor-num-nodes 1 \
  --actor-num-gpus-per-node 4 \
  --colocate \
  --rollout-num-gpus-per-engine 1 \
  --sglang-mem-fraction-static 0.5 \
  --sglang-tool-call-parser qwen25 \
  --use-session-server \
  --session-server-workers 8 \
  --tito-model qwen3 \
  --custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate \
  --custom-agent-function-path agent_env.training.miles.agent.run \
  --custom-rm-path agent_env.training.miles.agent.reward \
  --dynamic-sampling-filter-path agent_env.training.miles.agent.keep_scored_groups
EOF
} > train-cmd.sh
# Ray joins an entrypoint's words for a shell, so the command goes in a script the job runs.
ray job submit --address http://127.0.0.1:8265 \
  --runtime-env-json '{"env_vars": {"PYTHONPATH": "/root/miles:/root/Megatron-LM", "PYTHONUNBUFFERED": "1",
    "NCCL_NVLS_ENABLE": "0", "no_proxy": "127.0.0.1", "MASTER_ADDR": "127.0.0.1"}}' \
  -- bash "$PWD/train-cmd.sh"