# NeMo RL's GRPO on the GridSim v1 bundle, through agent_env.training.nemo_gym's resources server: the run on the
# AgentEnv docs' NeMo RL page, on NeMo RL main dd2e251 and four L40S. Paths are relative to the NeMo RL checkout.
# Run it with: uv run python examples/nemo_gym/run_grpo_nemo_gym.py --config grpo_gridsim.yaml
grpo:
  max_num_epochs: 1
  num_prompts_per_step: 4
  num_generations_per_prompt: 8
  max_rollout_turns: 1
  max_num_steps: 40
  normalize_rewards: false
  use_leave_one_out_baseline: false
  val_period: 20
  val_start_at: -1
  val_at_start: true
  val_at_end: true
  overlong_filtering: false
  advantage_clip_low: null
  advantage_clip_high: null
  max_val_samples: null
  val_num_generations_per_prompt: 1
  stop_at_validation_metric: null
  stop_at_validation_threshold: null
  val_batch_size: null
  seed: 42
  use_dynamic_sampling: false
  dynamic_sampling_max_gen_batches: 10
  batch_multiplier: 1
  reward_shaping:
    enabled: false
    overlong_buffer_length: 128
    overlong_buffer_penalty: 1
    max_response_length: ${policy.max_total_sequence_length}
  adv_estimator:
    name: grpo
    normalize_rewards: ${grpo.normalize_rewards}
    use_leave_one_out_baseline: ${grpo.use_leave_one_out_baseline}
    minus_baseline: true
  reward_scaling:
    enabled: false
    source_min: 0.0
    source_max: 1.0
    target_min: 0.0
    target_max: 1.0
  skip_reference_policy_logprobs_calculation: true
  seq_logprob_error_threshold: null
  invalid_tool_call_advantage: null
  malformed_thinking_advantage: null
  async_grpo:
    enabled: false
    max_trajectory_age_steps: 1
    max_generation_failures: 0
    in_flight_weight_updates: false
    recompute_kv_cache_after_weight_updates: false
loss_fn:
  reference_policy_kl_penalty: 0
  reference_policy_kl_type: k3
  kl_input_clamp_value: 20.0
  kl_output_clamp_value: 10.0
  ratio_clip_min: 0.2
  ratio_clip_max: 0.28
  ratio_clip_c: null
  use_on_policy_kl_approximation: false
  truncated_importance_sampling_ratio: null
  use_importance_sampling_correction: false
  token_level_loss: true
checkpointing:
  enabled: true
  checkpoint_dir: results/grpo-gridsim
  metric_name: val:accuracy
  higher_is_better: true
  keep_top_k: 2
  save_period: 10
  checkpoint_must_save_by: null
  save_optimizer: true
policy:
  model_name: qwen3-4b-gridsim-sft
  tokenizer:
    name: ${policy.model_name}
    chat_template_kwargs:
      enable_thinking: false
  hf_config_overrides: {}
  train_global_batch_size: 32
  train_micro_batch_size: 2
  logprob_batch_size: 2
  generation_batch_size: 32
  max_total_sequence_length: 3072
  precision: bfloat16
  logprob_chunk_size: null
  automodel_cfg:
    enabled: true
    checkpoint:
      model_save_format: safetensors
      save_consolidated: 'false'
      single_rank_consolidation: false
      consolidation_timeout_minutes: 30
    cpu_offload: false
    sequence_parallel: false
    activation_checkpointing: true
    deterministic: false
    tensor_parallel_size: 1
    context_parallel_size: 1
    custom_parallel_plan: null
    automodel_kwargs: {}
    lora_cfg:
      enabled: true
      target_modules: []
      exclude_modules: []
      match_all_linear: true
      dim: 32
      alpha: 64
      dropout: 0.0
      dropout_position: post
      lora_A_init: xavier
      use_triton: false
      restore_from: null
  megatron_cfg:
    enabled: false
    checkpoint:
      async_save: true
    bias_activation_fusion: false
    tensor_model_parallel_size: 2
    empty_unused_memory_level: 0
    activation_checkpointing: true
    train_iters: 100000
    expert_tensor_parallel_size: 1
    expert_model_parallel_size: 1
    pipeline_model_parallel_size: 1
    num_layers_in_first_pipeline_stage: null
    num_layers_in_last_pipeline_stage: null
    context_parallel_size: 1
    pipeline_dtype: ${policy.precision}
    sequence_parallel: false
    freeze_moe_router: true
    moe_router_dtype: fp64
    moe_router_load_balancing_type: none
    moe_router_bias_update_rate: 0.0
    apply_rope_fusion: true
    defer_fp32_logits: false
    moe_permute_fusion: true
    moe_enable_deepep: false
    moe_token_dispatcher_type: alltoall
    moe_shared_expert_overlap: false
    gradient_accumulation_fusion: false
    use_fused_weighted_squared_relu: false
    optimizer:
      optimizer: adam
      lr: 5.0e-06
      min_lr: 5.0e-07
      weight_decay: 0.01
      bf16: true
      fp16: false
      params_dtype: float32
      adam_beta1: 0.9
      adam_beta2: 0.999
      adam_eps: 1e-8
      sgd_momentum: 0.9
      use_distributed_optimizer: true
      use_precision_aware_optimizer: true
      optimizer_cpu_offload: false
      optimizer_offload_fraction: 0.0
      overlap_cpu_optimizer_d2h_h2d: false
      clip_grad: ${policy.max_grad_norm}
    scheduler:
      start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
      end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
      weight_decay_incr_style: constant
      lr_decay_style: constant
      lr_decay_iters: 100000
      lr_warmup_iters: 13
      lr_warmup_init: 5.0e-07
      override_opt_param_scheduler: true
    distributed_data_parallel_config:
      grad_reduce_in_fp32: false
      overlap_grad_reduce: true
      overlap_param_gather: true
      use_custom_fsdp: false
      data_parallel_sharding_strategy: optim_grads_params
    env_vars: null
  draft:
    enabled: false
    model_name: null
    loss_weight: 0.1
    speculator_type: eagle3
    num_layers: null
    aux_layer_indices: null
  dynamic_batching:
    enabled: false
    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
    sequence_length_round: 64
  sequence_packing:
    enabled: false
    train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}}
    logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}}
    algorithm: modified_first_fit_decreasing
    sequence_length_round: 64
  make_sequence_length_divisible_by: 1
  max_grad_norm: 1.0
  offload_optimizer_for_logprob: false
  offload_policy_before_refit: false
  offload_optimizer_for_refit: true
  optimizer:
    name: torch.optim.AdamW
    kwargs:
      lr: 2.0e-05
      weight_decay: 0.0
      betas:
      - 0.9
      - 0.999
      eps: 1.0e-08
  scheduler:
  - name: torch.optim.lr_scheduler.LinearLR
    kwargs:
      start_factor: 0.01
      end_factor: 1.0
      total_iters: 3
  - name: torch.optim.lr_scheduler.CosineAnnealingLR
    kwargs:
      T_max: 37
      eta_min: 0.0
  - milestones:
    - 3
  generation:
    backend: vllm
    max_new_tokens: 1024
    temperature: 1.1
    top_p: 1.0
    top_k: null
    val_temperature: ${.temperature}
    val_top_p: ${.top_p}
    val_top_k: ${.top_k}
    stop_token_ids: null
    stop_strings: null
    vllm_cfg:
      async_engine: true
      precision: ${policy.precision}
      refit_with_reload_api: false
      tensor_parallel_size: 1
      pipeline_parallel_size: 1
      expert_parallel_size: 1
      gpu_memory_utilization: 0.5
      max_model_len: 3072
      enforce_eager: false
      use_deep_gemm: false
      num_last_layers_in_bf16: 0
      num_first_layers_in_bf16: 0
      kv_cache_dtype: auto
      expose_http_server: true
      skip_tokenizer_init: false
      http_server_serving_chat_kwargs:
        enable_auto_tools: true
        tool_parser: hermes
    vllm_kwargs: {}
    colocated:
      enabled: true
      resources:
        gpus_per_node: null
        num_nodes: null
    lora_refit_mode: merged
data:
  max_input_seq_length: null
  shuffle: true
  num_workers: 0
  use_multiple_dataloader: false
  train:
    data_path: 3rdparty/Gym-workspace/Gym/resources_servers/gridsim_v1/data/train.jsonl
  validation:
    data_path: 3rdparty/Gym-workspace/Gym/resources_servers/gridsim_v1/data/eval.jsonl
  default:
    dataset_name: NemoGymDataset
    env_name: nemo_gym
    prompt_file: null
    system_prompt_file: null
    processor: nemo_gym_data_processor
env:
  should_use_nemo_gym: true
  should_log_nemo_gym_responses: true
  should_mask_flagged_samples: true
  nemo_gym:
    port_range_low: 5000
    port_range_high: 5999
    is_trajectory_collection: false
    config_paths:
    - responses_api_models/vllm_model/configs/vllm_model_for_training.yaml
    - resources_servers/gridsim_v1/configs/gridsim_v1.yaml
    policy_model:
      responses_api_models:
        vllm_model:
          uses_reasoning_parser: false
          extra_body:
            chat_template_kwargs:
              enable_thinking: false
    gridsim_v1_simple_agent:
      responses_api_agents:
        simple_agent:
          max_steps: 4
logger:
  log_dir: logs/grpo-gridsim
  num_val_samples_to_print: 0
  wandb_enabled: false
  tensorboard_enabled: true
  mlflow_enabled: false
  swanlab_enabled: false
  monitor_gpus: true
  wandb:
    project: grpo-gridsim
    name: qwen3-4b-gridsim
    log_nemo_gym_full_result_tables: false
  tensorboard: {}
  mlflow:
    experiment_name: grpo-gridsim
    run_name: qwen3-4b-gridsim
  gpu_monitoring:
    collection_interval: 10
    flush_interval: 10
cluster:
  gpus_per_node: 4
  num_nodes: 1
