cotrain-droid/config.yaml

output_dir: ./runs/cotrain_droid_teleop/2026-06-22_11-40-17
batch_size: 8
num_workers: 4
prefetch_factor: 4
lr_scheduler_type: constant
learning_rate: 1.0e-05
num_epochs: null
log_every: 10
save_every: 2000
save_accelerator_state: true
eval_every: 500
eval_num_inference_samples: 8
eval_num_inference_steps: 10
eval_video_fps: 16
gradient_accumulation_steps: 2
mixed_precision: bf16
seed: 42
max_grad_norm: 1.0
weight_decay: 0.001
resume: runs/cotrain_droid_teleop/2026-06-21_23-34-05/checkpoints/state/step_00006000
wandb:
  enabled: true
  workspace: null
  project: fastwam
  name: finetune
  group: null
  mode: online
data:
  train:
    _target_: fastwam.datasets.lerobot.lerobot_pretrain_dataset_v3.LerobotPretrainDatasetV3
    data_mix:
      droid: 0.5
      teleop: 0.5
    data_root_per_dataset:
      droid: ./data/datasets/lerobot_v3/droid
      teleop: data/datasets/finetune/pick_n_drop_202/train
    num_frames: 17
    action_video_freq_ratio: 2
    frame_size:
    - 256
    - 320
    camera_mapping_per_dataset:
      droid:
        primary: observation.images.exterior_image_1_left
        wrist: observation.images.wrist_image_left
      teleop:
        primary: exterior_image_1_left
        wrist: wrist_image_left
    action_mapping_per_dataset:
      droid:
        skip: 8
      teleop:
        actions: 8
    proprio_mapping_per_dataset:
      droid:
        skip: 8
      teleop:
        joint_position: 7
        gripper_position: 1
    action_norm_mask_per_dataset:
      droid:
        action:
        - true
        - true
        - true
        - true
        - true
        - true
        - false
        - false
      teleop:
        actions:
        - true
        - true
        - true
        - true
        - true
        - true
        - false
        - false
    proprio_norm_mask_per_dataset:
      droid:
        action:
        - true
        - true
        - true
        - true
        - true
        - true
        - false
        - false
      teleop:
        joint_position:
        - true
        - true
        - true
        - true
        - true
        - true
        - true
        gripper_position:
        - true
    emit_raw_video: true
    context_len: 128
    image_aug: false
    sampler_type: hierarchical_chunk
    episode_weighting: uniform
    sampler_num_samples: 1000000000
  val:
    _target_: fastwam.datasets.lerobot.lerobot_dataset_v3.LerobotDatasetV3
    data_root: data/datasets/finetune/pick_n_drop_202/val
    num_frames: 17
    action_video_freq_ratio: 2
    frame_size:
    - 256
    - 320
    camera_mapping:
      primary: exterior_image_1_left
      wrist: wrist_image_left
    action_mapping:
      actions: 8
    proprio_mapping:
      joint_position: 7
      gripper_position: 1
    action_norm_mask:
      actions:
      - true
      - true
      - true
      - true
      - true
      - true
      - false
      - false
    proprio_norm_mask:
      joint_position:
      - true
      - true
      - true
      - true
      - true
      - true
      - true
      gripper_position:
      - true
    emit_video: true
    context_len: 128
    image_aug: false
    sampler_type: resumable_epoch
    sampler_num_samples: 1000000000
model:
  _target_: fastwam.runtime.create_fastwam_joint
  model_id: Wan-AI/Wan2.2-TI2V-5B
  tokenizer_model_id: Wan-AI/Wan2.1-T2V-1.3B
  tokenizer_max_len: 128
  load_text_encoder: false
  redirect_common_files: false
  mot_checkpoint_mixed_attn: true
  action_dit_pretrained_path: checkpoints/ActionDiT_linear_interp_Wan22_alphascale_512hdim.pt
  skip_dit_load_from_pretrain: false
  num_frames: 17
  video_dit_config:
    patch_size:
    - 1
    - 2
    - 2
    in_dim: 48
    hidden_dim: 3072
    ffn_dim: 14336
    freq_dim: 256
    text_dim: 4096
    out_dim: 48
    num_heads: 24
    attn_head_dim: 128
    num_layers: 30
    eps: 1.0e-06
    use_gradient_checkpointing: true
  action_dit_config:
    proprio_dim: 8
    action_dim: 8
    hidden_dim: 512
    ffn_dim: 2048
    num_heads: 24
    attn_head_dim: 128
    num_layers: 30
    text_dim: 4096
    freq_dim: 256
    eps: 1.0e-06
    use_gradient_checkpointing: true
  video_scheduler:
    train_shift: 5.0
    infer_shift: 5.0
    num_train_timesteps: 1000
  action_scheduler:
    train_shift: 5.0
    infer_shift: 5.0
    num_train_timesteps: 1000
  loss:
    lambda_action: 1.0
name: cotrain
max_steps: 20000