# PACE-SimplerEnv: model config of the released PACE checkpoint, read by the PACE code together with pytorch_model.pt
# and dataset_statistics.json in this directory:
#   from pace.model.framework.base_framework import baseframework
#   policy = baseframework.from_pretrained("./PACE-SimplerEnv")
# Training recipes and launchers are in the code repository (docs/TRAINING.md).
run_id: PACE-SimplerEnv
run_root_dir: runs
seed: 42
trackers: [jsonl]  # "jsonl" -> <run>/metrics.jsonl; add "wandb" for Weights & Biases (entity: wandb_entity or $WANDB_ENTITY)
wandb_entity: null
wandb_project: pace

framework:
  # "PACE"; the long alias "ProgressVLA_DinoTargetQueryProgressTripleReadoutFullGoal" is the name used by the
  # original training runs and is still accepted.
  name: PACE
  qwenvl:
    base_vlm: Qwen/Qwen3-VL-4B-Instruct  # env BASE_VLM (local directory or HF id) takes precedence
    attn_implementation: flash_attention_2  # env PACE_ATTN_IMPLEMENTATION (e.g. sdpa) overrides
  action_model:
    action_model_type: DiT-B  # 12 heads x 64, 16 layers (diffusion_model_cfg.num_layers)
    hidden_size: 1024
    add_pos_embed: true
    max_seq_len: 1024
    action_dim: 7
    state_dim: 7  # the state encoder exists in the checkpoint, but no proprioceptive state is fed
    future_action_window_size: 15
    action_horizon: 16  # action chunk H = future_action_window_size + 1
    past_action_window_size: 0
    noise_beta_alpha: 1.5
    noise_beta_beta: 1.0
    noise_s: 0.999
    num_timestep_buckets: 1000
    num_inference_timesteps: 4  # flow-matching (Euler) denoising steps
    num_target_vision_tokens: 32
    diffusion_model_cfg:
      dropout: 0.2
      final_dropout: true
      interleave_self_attention: true
      norm_type: ada_norm
      num_layers: 16
      output_dim: 1024
      positional_embeddings: null
  progress:
    layer_idx: 18  # query readouts from the hidden states of VLM layer 18
    num_goal_queries: 4  # K goal queries, mean-pooled into the foresight-anchor head
    use_goal_dit_token: true  # the predicted anchor is a DiT conditioning token (besides p_cur and dp_cur)
    lambda_long: 0.1  # goal-anchor loss (MSE in DINOv2 space)
    lambda_pcur: 0.1  # current-progress loss
    lambda_dpcur: 0.1  # forward-commitment loss
    lambda_cons: 0.1  # consistency: p_cur readout vs projection onto the predicted anchor
    cons_detach_anchor: true
    query_init_std: 0.01
    train_query_token_rows_only: true  # new token rows are trained; the rest of the VLM is fully fine-tuned
    progress_future_offset: 15  # Delta (history / commitment horizon in env steps); equals datasets.vla_data value
    use_history_images: true  # the frame at t - Delta is an extra VLM input
    use_dp_prev_input: true  # the historical pace dp_prev enters the VLM as a token
    use_full_goal_anchor_at_inference: true
    infer_history_mode: step_aligned
    dino_variant: vitb14  # frozen DINOv2 ViT-B/14 (progress geometry)
    dino_target_dim: 768
    dino_head_hidden_dim: 1024
    dino_repo_root: null  # null -> vendored third_party/dinov2 (env DINO_REPO_ROOT overrides)
    dino_weights: null  # null -> $PACE_CACHE/dinov2 download (env DINOV2_WEIGHTS overrides)
    # Commitment audit (inference-time failure detector, never changes actions): alarm when the shortfall CUSUM
    # exceeds audit_threshold or p_cur has not risen for more than audit_stall_chunks action chunks.
    # audit_threshold is shared by all benchmarks (unmet commitments adding up to 0.4 of the task); the stall limit is
    # calibrated per benchmark on held-out checkpoints (2.5% false alarms on successful episodes); see docs/PROGRESS.md.
    audit_kappa: 0.0
    audit_threshold: 0.4
    audit_stall_chunks: 2.12

datasets:
  vla_data:
    dataset_py: lerobot_datasets
    image_size: [224, 224]  # inference-time resize of the client images
    CoT_prompt: "Your task is {instruction}. To identify the key objects for your task. Locate their bounding boxes in [x1,y1,x2,y2] format."
    return_progress_anchors: true
    progress_future_offset: 15

trainer:  # inference only reads these
  repeated_diffusion_steps: 4
  pretrained_checkpoint: null