qwen_gr00t.yaml

system:
  batch_size: 8
  train_steps: 30000
  log_freq: 1
  grad_clip_norm: 1.0
  use_amp: true
  shuffle: true
  num_workers: 4
  # Weight applied to VLM language modelling loss when co-training with vlm_data_path.
  # Set to 0 or omit vlm_data_path to disable co-training.
  vlm_loss_scale: 0.1

  checkpoint:
    output_directory: ${experiment.exp_dir}
    # Whether to save checkpoint
    save_checkpoint: true
    # Number of steps between checkpoints
    save_freq: 10000
    # Path to a checkpoint directory to resume training from (e.g. /path/to/checkpoints/005000)
    # resume_from:

  # --- Activation Checkpointing ---
  # Modes: "none" (default), "full", "selective", "memory_budget"
  # "full": wraps every FSDP unit with gradient checkpointing
  # "selective": wraps every Nth layer (selective_ac_option="2" = every 2nd)
  #              or uses per-op SAC (selective_ac_option="op")
  # "memory_budget": compiler-driven (requires torch.compile)
  activation_checkpoint:
    mode: none
    # For selective mode: "2" = every 2nd layer, "3" = every 3rd, "op" = per-op SAC
    # selective_ac_option: "2"
    preserve_rng_state: false
    # For memory_budget mode only:
    # memory_budget: 0.5
    # Only checkpoint the action model (DiT blocks). VLM layers use dynamic
    # sequence lengths (variable vision tokens) which breaks recomputation.
    # checkpoint_patterns:
    #   - "action_model\\._head\\.model\\.transformer_blocks\\.\\d+$"

model:
  model_name: qwen_gr00t
  vlm:
    type: qwen3-vl
    base_vlm: /workspace/models/Qwen/Qwen3-VL-4B-Instruct/
    attn_implementation: flash_attention_2
  action_model:
    # Whether to condition the action model on proprioceptive state (observation.state)
    use_state: false
    type: gr00t_action_head
    action_model_type: DiT-B
    hidden_size: 1024
    add_pos_embed: True
    max_seq_len: 1024
    action_dim: 7
    state_dim: 7
    future_action_window_size: 7
    action_horizon: 8
    repeated_diffusion_steps: 4
    noise_beta_alpha: 1.5
    noise_beta_beta: 1.0
    noise_s: 0.999
    num_timestep_buckets: 1000
    num_inference_timesteps: 4
    num_target_vision_tokens: 32
    diffusion_model_cfg:
      cross_attention_dim: 2048
      dropout: 0.2
      final_dropout: True
      interleave_self_attention: True
      norm_type: ada_norm
      num_layers: 16
      output_dim: 1024
      positional_embeddings: None

  prompt_template: "Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format."

  # --- Chunked Cross Entropy ---
  # 0 = disabled
  # >0 = chunk size in tokens. Processes logits in fixed-size chunks to reduce
  #       peak memory from the [batch*seq_len, vocab_size] float32 softmax tensor.
  chunked_ce_tokens: 0

  normalization_mapping:
    VISUAL: IDENTITY
    STATE: MIN_MAX
    ACTION: MIN_MAX

  optimizer:
    name: AdamW
    lr: 2.5e-5
    betas: [0.9, 0.95]
    eps: 1.0e-08
    weight_decay: 1.0e-08
    param_groups:
      vlm:
        lr: 1.0e-05
      action_model:
        lr: 1.0e-04
    scheduler:
      name: cosine_with_min_lr
      warmup_steps: 5000
      scheduler_kwargs:
        min_lr: 1.0e-06
      # Legacy fields kept for BC
      decay_steps: 30000
      decay_lr: 2.5e-6

  # ============================================================
  # Module Freezing Configuration
  # ============================================================
  # Freezing logic: freeze_patterns are applied first, then keep_patterns override.
  # Patterns are regex matched against full parameter names.
  #
  # Common patterns for QwenGR00T:
  #   - "qwen_vl_interface\\..*"                              # Entire VLM
  #   - "qwen_vl_interface\\.model\\.visual\\..*"             # Vision encoder
  #   - "qwen_vl_interface\\.model\\.model\\..*"              # Language model
  #   - "qwen_vl_interface\\.model\\.model\\.layers\\.[0-9]\\."  # LLM layers 0-9
  #   - "action_model\\..*"                                   # Action head
  #   - "action_model\\.model\\.transformer_blocks\\.[0-7]\\."   # DiT blocks 0-7
  #
  # freeze:
  #   # SCENARIO A: Freeze VLM, train only action head
    # freeze_patterns:
    #   - "qwen_vl_interface\\..*"
  #
  #   # SCENARIO B: Freeze VLM but keep projector trainable
  #   # freeze_patterns:
  #   #   - "qwen_vl_interface\\..*"
  #   # keep_patterns:
  #   #   - "qwen_vl_interface\\.model\\.visual\\.merger\\..*"
  #
  #   # SCENARIO C: Freeze everything except action decoder
  #   # freeze_patterns:
  #   #   - ".*"
  #   # keep_patterns:
  #   #   - "action_model\\.action_decoder\\..*"

data:
  dataset_type: lerobot
  wds:
    vision_root: ""
    action_key: eepose
    state_key: eepose
  # Path to the training data
  data_path: /workspace/datasets/IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot/
  # Path to VLM co-training data (WDS/Energon format). Leave unset to disable co-training.
  # vlm_data_path: /workspace/datasets/vlm_cotrain/
  tolerance_s: 0.0001

  # --- Data Mixture ---
  # When data_mix is set, it takes precedence over data_path.
  # Each entry specifies a dataset path and sampling weight.
  # Weights are relative — they get normalized internally.
  # Optional: data_root_dir as a common prefix for relative paths.
  #
  # data_root_dir: /workspace/datasets/
  # data_mix:
  #   - path: IPEC-COMMUNITY/libero_goal_no_noops_1.0.0_lerobot
  #     weight: 1.0
  #   - path: IPEC-COMMUNITY/libero_spatial_no_noops_1.0.0_lerobot
  #     weight: 0.5
  # balance_dataset_weights: true  # scale weights by dataset size (default: true)

  preprocessor:
    name: policy_preprocessor
    steps:
      - registry_name: rename_observations_processor
        config:
          rename_map: {}
      - registry_name: to_batch_processor
        config: {}
      - registry_name: device_processor
        config:
          device: cuda
          float_dtype: null
      - registry_name: normalizer_processor
        config:
          eps: 1e-8
          features: {}
          # norm_map is injected at runtime from model.normalization_mapping
  postprocessor:
    name: policy_postprocessor
    steps:
      - registry_name: unnormalizer_processor
        config:
          eps: 1e-8
          features: {}
          # norm_map is injected at runtime from model.normalization_mapping