Multi-Turn GRPO Text2SQL with Async Rollouts

You are a data science expert. Below, you are provided with a database schema and a natural language question. Your task is to understand the schema and generate a valid SQL query to answer the question within limited turns. You should breakdown the problem, draft your reasoning process, and generate the solution.

Database Engine:
SQLite

Database Schema:
{db_details}
This schema describes the database's structure, including tables, columns, primary keys, foreign keys, and any relevant relationships or constraints.

External Knowledge:
{external_knowledge}

Question:
{question}

Instructions:
- Make sure you only output the information that is asked in the question. If the question asks for a specific column, make sure to only include that column in the SELECT clause, nothing more.
- The generated query should return all of the information asked in the question without any missing or extra information.
- Before generating the final SQL query, please think through the steps of how to write the query. It should include detailed considerations such as analyzing questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, thinking of how to call SQL tools, and revisiting previous steps.

Format:
- Conduct thinking inside <think>...</think> blocks every time you get new observation or information. 
- You can use SQL tool written within a single <sql>your sql</sql> block to explore or verify. SQL tool output will be shown as dataframe inside <observation>...</observation>. Based on this observation, you can think again and refine.
- The returned dataframe will be truncated in 50 rows if observation is too long. 
- If you find no further exploration is needed or reaches max turns, you MUST directly provide the final SQL query solution inside <solution>...</solution>. 
...
{db_details}: 
CREATE TABLE country_aliases (
    alias_id INTEGER,
    country_id INTEGER,
    alias_name TEXT,
    alias_type TEXT,
    PRIMARY KEY (alias_id),
    FOREIGN KEY (country_id) REFERENCES countries (country_id)
);
...
{question}: 
Could you please provide me with the English name, current population, average population over the last five years, and population from last year for each country, based on the country's population metrics? I need this information to analyze population trends.

hf download NovaSky-AI/SkyRL-SQL-653-data-newfmt --local-dir $HOME/data/sql --repo-type dataset
hf download seeklhy/OmniSQL-datasets data.zip --repo-type dataset --local-dir <path_to_file.zip>
unzip <path_to_file.zip>

# path for dataset (.parquet files) containing the prompts and metadata for each question
DATA_DIR="$HOME/data/sql"
# path for .db files for environment interaction
DB_PATH="$HOME/path/to/db_files"

uv run --isolated --extra vllm -m skyrl_train.entrypoints.main_base \
    #### Environment configuration
    environment.env_class=text2sql \
    environment.skyrl_gym.text2sql.db_path=$DB_PATH \

    #### Multi-turn Async Rollouts configuration
    # this is used to set the max turns for the environment
    generator.max_turns=6 \
    # we need to make sure to set async_engine=true for async rollouts
    generator.async_engine=true \
    # we need to make sure to set batched=false for async rollouts
    generator.batched=false \

    #### context length related configurations
    # trainer.max_prompt_length is the max length of the initial prompt
    trainer.max_prompt_length=6000 \
    # generator.max_input_length is the max length of the input to the model after any number of turns (including the initial prompt)
    generator.max_input_length=29000 \
    # generator.sampling_params.max_generate_length is the max length of the generated response for EACH turn
    generator.sampling_params.max_generate_length=3000 \

    #### multi-turn generation format - see `skyrl_train/generators/skyrl_gym_generator.py` for more details
    generator.use_conversation_multi_turn=false \

    #### data configuration
    data.train_data="['$DATA_DIR/train.parquet']" \
    data.val_data="['$DATA_DIR/validation.parquet']" \

    #### Placement configuration - note since we set use_kl_loss=false below, we don't need to use a ref model
    # policy placement configuration
    trainer.policy.model.path="Qwen/Qwen2.5-Coder-7B-Instruct" \
    trainer.placement.colocate_all=true \
    trainer.placement.policy_num_gpus_per_node=8 \
    # inference engine placement configuration
    generator.num_inference_engines=2 \
    generator.inference_engine_tensor_parallel_size=4 \

    #### algorithm configuration
    trainer.epochs=30 \
    trainer.algorithm.advantage_estimator="grpo" \
    trainer.algorithm.use_kl_loss=false \
    trainer.algorithm.use_kl_loss=false \
    generator.n_samples_per_prompt=5 \

    #### generation sampling params (relevant to algorithm correctness)
    generator.sampling_params.temperature=0.6 \
    generator.sampling_params.top_p=0.95 \
    generator.sampling_params.stop='["</sql>", "</solution>"]' \
    generator.eval_sampling_params.stop='["</sql>", "</solution>"]' \

    #### training configuration
    trainer.policy.optimizer_config.lr=1.0e-6 \
    trainer.train_batch_size=256 \
    trainer.policy_mini_batch_size=256 \
    trainer.micro_forward_batch_size_per_gpu=8 \
    trainer.micro_train_batch_size_per_gpu=1 \
    trainer.eval_batch_size=1024 \
    trainer.eval_before_train=true \
    trainer.eval_interval=5 \
    ... # Other parameters (see `examples/text_to_sql/run_skyrl_sql.sh` for the full script)

export WANDB_API_KEY=your_wandb_api_key
bash examples/text_to_sql/run_skyrl_sql.sh