Hugging Face's logo

OpenWAM
/
pretrain_robot_only

Robotics
openwam
world-action-model
Model card Files
xet
 Community
pretrain_robot_only
File size: 4,555 Bytes
86830a1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
model:
  video_backbone:
    encoder:
      name: wan22_vae
      model_path: /path/to/Wan2.2-TI2V-5B
    name: wan22_ti2v_5b
    model_path: /path/to/Wan2.2-TI2V-5B
    from_scratch: false
    shift_video: 5.0
    components:
    - attr: dit
      model_class: openwam.model.video_backbone.wan.models.dit.WanModel
      extra_kwargs:
        has_image_input: false
        patch_size:
        - 1
        - 2
        - 2
        in_dim: 48
        dim: 3072
        ffn_dim: 14336
        freq_dim: 256
        text_dim: 4096
        out_dim: 48
        num_heads: 24
        num_layers: 30
        eps: 1.0e-06
        seperated_timestep: true
        require_clip_embedding: false
        require_vae_embedding: false
        fuse_vae_embedding_in_latents: true
    - attr: vae
      model_class: openwam.model.video_backbone.wan.models.vae.WanVideoVAE38
      extra_kwargs: {}
    - attr: text_encoder
      model_class: openwam.model.video_backbone.wan.models.text_encoder.WanTextEncoder
      extra_kwargs: {}
    tokenizer:
      class: openwam.model.video_backbone.wan.models.text_encoder.HuggingfaceTokenizer
      attr: tokenizer
      subdir: tokenizer/google/umt5-xxl
      path_kwarg: name
      kwargs:
        seq_len: 512
        clean: whitespace
  action_backbone:
    dim: 1024
    ffn_dim: 4096
    shift_action: 5.0
  freeze:
  - video_backbone.text_encoder
  - video_backbone.reason1
  - video_backbone.vae
  - video_backbone.image_encoder
  - video_backbone.video_encoder
  architecture:
    framework: dual_system
    variant: joint_self_attn
    action_dim: 80
    use_proprioception: true
    state_dim: 80
    bridge_layers: null
    bridge_interval: 1
    mot_checkpoint_mixed_attn: true
    attention_mask_mode: action_sees_video
    video_attention_mask_mode: first_frame_causal
    detach_bridge: false
    idm_video_cond_noise_prob: 0.5
dataloader:
  datasets:
    robocoin:
      type: robocoin
      dataset_dir: /path/to/pretrain_dataset/RoboCOIN
      num_frames: 33
      video_stride: 4
      window_stride: 1
      height: 384
      width: 320
      multiview: true
      split: train
      color_jitter:
        brightness: 0.2
        contrast: 0.2
        saturation: 0.2
        hue: 0.0
      normalize_mode: quantile
      enable_action_supervision: true
      unify_action: true
      unify_action_map:
      - 0-9
      - 34-43
      total_hours: 600.0
      enabled: true
      weight: 1.0
    oxe_droid:
      type: oxe_droid
      dataset_dir: /path/to/pretrain_dataset/DROID-Dataset
      num_frames: 33
      video_stride: 4
      window_stride: 1
      height: 384
      width: 320
      multiview: true
      split: train
      color_jitter:
        brightness: 0.2
        contrast: 0.2
        saturation: 0.2
        hue: 0.0
      normalize_mode: quantile
      enable_action_supervision: true
      unify_action: true
      unify_action_map:
      - 0-9
      - 34-43
      total_hours: null
      enabled: false
      weight: 1.0
    egodex:
      type: egodex
      dataset_dir: /path/to/pretrain_dataset/EgoDex
      num_frames: 33
      video_stride: 4
      window_stride: 1
      height: 384
      width: 320
      target_camera: observation.images.ego
      multiview: true
      camera_layout:
      - observation.images.ego
      - __missing_left__
      - __missing_right__
      split: train
      color_jitter:
        brightness: 0.2
        contrast: 0.2
        saturation: 0.2
        hue: 0.0
      unify_action: true
      total_hours: 350.0
      enabled: false
      weight: 1.0
  type: mixture
  weight_strategy: proportional
  seed: 42
training:
  debug: false
  learning_rate: 0.0001
  adam_betas:
  - 0.9
  - 0.95
  weight_decay: 0.01
  max_grad_norm: 1.0
  num_epochs: 1
  max_steps: null
  batch_size: 128
  gradient_accumulation_steps: 1
  lr_scheduler: cosine
  warmup_ratio: 0.05
  lr_min_ratio: 0.01
  action_lr: null
  video_lr: null
  mixed_precision: bf16
  zero_stage: 2
  use_gradient_checkpointing: true
  use_gradient_checkpointing_offload: false
  initialize_model_on_cpu: false
  offload_optimizer_device: none
  lambda_video: 1.0
  lambda_action: 1.0
  max_timestep_boundary: 1.0
  min_timestep_boundary: 0.0
  output_path: outputs/openwam_checkpoints
  save_steps: 500
  save_full_states_for_resume: false
  keep_last_k_ckpts: 1
  finetune_ckpt_path: null
  resume_ckpt_path: null
  dataset_num_workers: 8
project:
  name: openwam
  output_dir: outputs/${now:%Y-%m-%d_%H-%M-%S}
  seed: 42
  wandb:
    project: openwam
    run_name: null
    entity: null