#!/usr/bin/env bash
# GRPO on the GridSim v1 bundle through AgentEnv's env and generator, as the run the SkyRL page plots ran it on one
# L40S: the Qwen3-4B SFT warm start, LoRA r32/alpha 64, lr 2e-5 cosine with 3 warmup steps, 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, 20 and 40. Run it from a folder holding qwen3-4b-gridsim-sft and the rows
# python -m agent_env.training.skyrl wrote to data/, in SkyRL's environment with agentenv-framework[skyrl] installed;
# with a server that asks for a token, export the variable the rows' --token-var names first.
set -euo pipefail
python -m agent_env.training.skyrl.train \
  data.train_data="['data/train.parquet']" \
  data.val_data="['data/eval.parquet']" \
  data.dataloader.num_workers=0 \
  trainer.algorithm.advantage_estimator=grpo \
  trainer.algorithm.use_kl_loss=false \
  trainer.policy.model.path=qwen3-4b-gridsim-sft \
  trainer.policy.model.lora.rank=32 \
  trainer.policy.model.lora.alpha=64 \
  trainer.policy.optimizer_config.lr=2e-5 \
  trainer.policy.optimizer_config.scheduler=cosine \
  trainer.policy.optimizer_config.num_warmup_steps=3 \
  trainer.strategy=fsdp \
  trainer.placement.colocate_all=true \
  trainer.placement.policy_num_gpus_per_node=1 \
  trainer.placement.ref_num_gpus_per_node=1 \
  generator.inference_engine.num_engines=1 \
  generator.inference_engine.tensor_parallel_size=1 \
  generator.inference_engine.backend=vllm \
  generator.inference_engine.run_engines_locally=true \
  generator.inference_engine.weight_sync_backend=nccl \
  generator.inference_engine.gpu_memory_utilization=0.5 \
  trainer.train_batch_size=4 \
  trainer.policy_mini_batch_size=4 \
  generator.n_samples_per_prompt=8 \
  trainer.micro_forward_batch_size_per_gpu=4 \
  trainer.micro_train_batch_size_per_gpu=4 \
  trainer.max_prompt_length=3072 \
  generator.max_input_length=8192 \
  generator.sampling_params.max_generate_length=1024 \
  generator.sampling_params.temperature=1.1 \
  generator.eval_sampling_params.temperature=1.1 \
  generator.eval_sampling_params.max_generate_length=1024 \
  generator.eval_n_samples_per_prompt=1 \
  trainer.eval_batch_size=200 \
  trainer.eval_before_train=true \
  trainer.eval_interval=20 \
  trainer.max_training_steps=40 \
  generator.max_turns=4 \
  generator.batched=false \
  generator.chat_template_kwargs.enable_thinking=false \
  environment.env_class=agentenv \
  environment.skyrl_gym.max_env_workers=64 \
  trainer.ckpt_interval=100000 \
  trainer.resume_mode=null \
  trainer.logger=console \
  trainer.project_name=agentenv \
  trainer.run_name=gridsim-skyrl \
  trainer.ckpt_path=ckpt \
  trainer.export_path=export