Upload folder using huggingface_hub

action_decoder/chain/w2a_so101_task2_combined_lora32_iter2500_75m_2gpu_bsz2_accum8_val5/config.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:326b78a2cd08812f0b3f13b46103b4181a7b9ec389c64502b269dac69a43ae84
+size 76376

action_decoder/chain/w2a_so101_task2_combined_lora32_iter2500_75m_2gpu_bsz2_accum8_val5/config.yaml

@@ -0,0 +1,531 @@
+checkpoint:
+  broadcast_via_filesystem: 'False'
+  dcp_allow_mismatched_size: 'False'
+  dcp_async_mode_enabled: 'False'
+  jit:
+    device: cuda
+    dtype: bfloat16
+    enabled: 'False'
+    input_shape: null
+    strict: 'True'
+  keys_not_to_resume: []
+  load_ema_to_reg: 'False'
+  load_path: ''
+  load_training_state: 'False'
+  only_load_scheduler_state: 'False'
+  save_iter: '100'
+  strict_resume: 'True'
+  type:
+    _target_: <class 'cosmos_predict2.checkpointer.Checkpointer'>
+    callbacks: null
+  verbose: 'True'
+data_config:
+  _target_: <function get_data_config at 0x77eab678a680>
+  config_name: so101_lerobot
+dataloader_train:
+  _target_: <class 'torch.utils.data.dataloader.DataLoader'>
+  batch_sampler: null
+  batch_size: '2'
+  collate_fn: null
+  dataset:
+    _target_: <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>
+    action_image_horizon: '56'
+    action_key: action
+    action_lowdim_horizon: '90'
+    action_shift_s: '0.2'
+    cache_dir: null
+    image_key: observation.images.front
+    image_resize:
+    - '480'
+    - '640'
+    language_embeddings_path: null
+    lowdim_target_fps: '30.0'
+    num_val_episodes: '1'
+    obs_image_horizon: '5'
+    repo_id: push-that-thing/task_2
+    root: null
+    seed: '42'
+    state_key: observation.state
+    t5_encoder_ckpt: null
+    target_fps: '5.0'
+    train: 'True'
+    val_ratio: '0.05'
+    video_backend: pyav
+  drop_last: 'True'
+  generator: null
+  in_order: 'True'
+  multiprocessing_context: null
+  num_workers: '8'
+  persistent_workers: 'True'
+  pin_memory: 'True'
+  pin_memory_device: ''
+  prefetch_factor: '8'
+  sampler:
+    _target_: <class 'cosmos_predict2.data.resumable_sampler.ResumableDistributedSampler'>
+    dataset:
+      _target_: <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>
+      action_image_horizon: '56'
+      action_key: action
+      action_lowdim_horizon: '90'
+      action_shift_s: '0.2'
+      cache_dir: null
+      image_key: observation.images.front
+      image_resize:
+      - '480'
+      - '640'
+      language_embeddings_path: null
+      lowdim_target_fps: '30.0'
+      num_val_episodes: '1'
+      obs_image_horizon: '5'
+      repo_id: push-that-thing/task_1
+      root: null
+      seed: '42'
+      state_key: observation.state
+      t5_encoder_ckpt: null
+      target_fps: '5.0'
+      train: 'True'
+      val_ratio: null
+      video_backend: pyav
+    num_replicas:
+      _target_: <function get_data_parallel_world_size at 0x77ed19e6d5a0>
+      partial_data_parallel: 'False'
+      with_context_parallel: 'False'
+    rank:
+      _target_: <function get_data_parallel_rank at 0x77ed19e6d6c0>
+      partial_data_parallel: 'False'
+      with_context_parallel: 'False'
+    seed: '0'
+    shuffle: 'True'
+  shuffle: null
+  timeout: '0'
+  worker_init_fn: null
+dataloader_val:
+  _target_: <class 'torch.utils.data.dataloader.DataLoader'>
+  batch_sampler: null
+  batch_size: '1'
+  collate_fn: null
+  dataset:
+    _target_: <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>
+    action_image_horizon: '56'
+    action_key: action
+    action_lowdim_horizon: '90'
+    action_shift_s: '0.2'
+    cache_dir: null
+    image_key: observation.images.front
+    image_resize:
+    - '480'
+    - '640'
+    language_embeddings_path: null
+    lowdim_target_fps: '30.0'
+    num_val_episodes: '1'
+    obs_image_horizon: '5'
+    repo_id: push-that-thing/task_2
+    root: null
+    seed: '42'
+    state_key: observation.state
+    t5_encoder_ckpt: null
+    target_fps: '5.0'
+    train: 'False'
+    val_ratio: '0.05'
+    video_backend: pyav
+  drop_last: 'False'
+  generator: null
+  in_order: 'True'
+  multiprocessing_context: null
+  num_workers: '0'
+  persistent_workers: 'False'
+  pin_memory: 'False'
+  pin_memory_device: ''
+  prefetch_factor: null
+  sampler:
+    _target_: <class 'cosmos_predict2.data.resumable_sampler.ResumableDistributedSampler'>
+    dataset:
+      _target_: <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>
+      action_image_horizon: '56'
+      action_key: action
+      action_lowdim_horizon: '90'
+      action_shift_s: '0.2'
+      cache_dir: null
+      image_key: observation.images.front
+      image_resize:
+      - '480'
+      - '640'
+      language_embeddings_path: null
+      lowdim_target_fps: '30.0'
+      num_val_episodes: '1'
+      obs_image_horizon: '5'
+      repo_id: push-that-thing/task_1
+      root: null
+      seed: '42'
+      state_key: observation.state
+      t5_encoder_ckpt: null
+      target_fps: '5.0'
+      train: 'False'
+      val_ratio: null
+      video_backend: pyav
+    num_replicas:
+      _target_: <function get_data_parallel_world_size at 0x77ed19e6d5a0>
+      partial_data_parallel: 'False'
+      with_context_parallel: 'False'
+    rank:
+      _target_: <function get_data_parallel_rank at 0x77ed19e6d6c0>
+      partial_data_parallel: 'False'
+      with_context_parallel: 'False'
+    seed: '0'
+    shuffle: 'False'
+  shuffle: null
+  timeout: '0'
+  worker_init_fn: null
+defaults:
+- _self_
+- data_config: null
+- video_dataset_train: null
+- video_dataset_val: null
+- dataloader_train: null
+- dataloader_val: null
+- world2action_pipe: null
+- optimizer: fusedadamw
+- scheduler: constant
+- model: null
+- callbacks:
+  - basic
+- net: null
+- ema: null
+- checkpoint: null
+- ckpt_type: null
+- experiment: null
+job:
+  group: so101
+  name: w2a_so101_task2_combined_lora32_iter2500_75m_2gpu_bsz2_accum8_val5
+  project: vam
+model:
+  _target_: <class 'cosmos_predict2.models.world2action_model.World2ActionModel'>
+  config:
+    _target_: cosmos_predict2.models.world2action_model.World2ActionModelConfig
+    action_dit_path: ''
+    data_config:
+      _target_: <function get_data_config at 0x77eab678a680>
+      config_name: so101_lerobot
+    ema:
+      enabled: false
+      iteration_shift: 0
+      rate: 0.1
+    fsdp_shard_size: '0'
+    init_lora_weights: 'True'
+    lora_alpha: '16'
+    lora_rank: '16'
+    lora_target_modules: q_proj,k_proj,v_proj,output_proj,mlp.layer1,mlp.layer2
+    loss_reduce: mean
+    loss_scale: '10.0'
+    pipe_config:
+      _target_: cosmos_predict2.configs.config_world2action.World2ActionPipelineConfig
+      ema:
+        enabled: false
+        iteration_shift: 0
+        rate: 0.1
+      net:
+        _target_: <class 'cosmos_predict2.models.world2action_dit.World2ActionDIT'>
+        adaln_lora_dim: '128'
+        atten_backend: flash_attn_no_cp
+        crossattn_emb_channels: '2048'
+        in_channels: '6'
+        max_horizon: '91'
+        mlp_ratio: '4.0'
+        model_channels: '512'
+        num_blocks: '12'
+        num_heads: '8'
+        out_channels: '6'
+        pair_timestep_feature_rank: '512'
+        sac_config:
+          every_n_blocks: 1
+          mode: none
+        use_adaln_lora: 'True'
+      precision: bfloat16
+      scheduler:
+        alpha: 1.0
+        beta: 1.0
+        num_denoising_steps: 10
+      xattn_layer_idx: '20'
+    precision: bfloat16
+    train_architecture: base
+    validation_mode: loss_only
+    video_dit_path: /home/nvidia/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt
+    video_pipe_config:
+      adjust_video_noise: true
+      conditioner:
+        _target_: <class 'cosmos_predict2.conditioner.VideoConditioner'>
+        fps:
+          _target_: <class 'cosmos_predict2.conditioner.ReMapkey'>
+          dropout_rate: '0.0'
+          dtype: null
+          input_key: fps
+          output_key: fps
+        padding_mask:
+          _target_: <class 'cosmos_predict2.conditioner.ReMapkey'>
+          dropout_rate: '0.0'
+          dtype: null
+          input_key: padding_mask
+          output_key: padding_mask
+        text:
+          _target_: <class 'cosmos_predict2.conditioner.TextAttr'>
+          dropout_rate: '0.0'
+          input_key:
+          - obs/language_embedding
+        use_video_condition:
+          _target_: <class 'cosmos_predict2.conditioner.BooleanFlag'>
+          dropout_rate: '0.0'
+          input_key: fps
+          output_key: use_video_condition
+      conditioning_strategy: frame_replace
+      ema:
+        _target_: cosmos_predict2.configs.defaults.ema.EMAConfig
+        enabled: 'False'
+        iteration_shift: '0'
+        rate: '0.1'
+      guardrail_config:
+        checkpoint_dir: /home/nvidia/pdt-mimic/mimic-video/model/checkpoints
+        enabled: false
+        offload_model_to_cpu: true
+      input_image_key: images
+      input_video_key: video
+      max_num_conditional_frames: 2
+      min_num_conditional_frames: 1
+      net:
+        _target_: <class 'cosmos_predict2.models.video2world_dit.MinimalV1LVGDiT'>
+        adaln_lora_dim: '256'
+        atten_backend: minimal_a2a
+        concat_padding_mask: 'True'
+        extra_per_block_abs_pos_emb: 'False'
+        in_channels: '16'
+        max_frames: '128'
+        max_img_h: '240'
+        max_img_w: '240'
+        model_channels: '2048'
+        num_blocks: '28'
+        num_heads: '16'
+        out_channels: '16'
+        patch_spatial: '2'
+        patch_temporal: '1'
+        pos_emb_cls: rope3d
+        pos_emb_interpolation: crop
+        pos_emb_learnable: 'True'
+        rope_enable_fps_modulation: 'False'
+        rope_h_extrapolation_ratio: '3.0'
+        rope_t_extrapolation_ratio: '1.0'
+        rope_w_extrapolation_ratio: '3.0'
+        sac_config:
+          _target_: cosmos_predict2.models.text2image_dit.SACConfig
+          every_n_blocks: '1'
+          mode: predict2_2b_720
+        use_adaln_lora: 'True'
+      precision: bfloat16
+      rectified_flow_loss_weight_uniform: true
+      rectified_flow_t_scaling_factor: 1.0
+      resize_online: false
+      resolution: '480'
+      sigma_conditional: 0.0001
+      sigma_data: 1.0
+      state_ch: 16
+      state_t: 16
+      text_encoder:
+        cls: !!python/object/apply:imaginaire.constants.TextEncoderClass
+        - t5
+        t5:
+          ckpt_path: /home/nvidia/pdt-mimic/mimic-video/model/checkpoints/text_encoder/t5-11b
+          embed_dim: 1024
+          num_tokens: 512
+      timestamps:
+        is_forward: false
+        nfe: 35
+        order: 7.0
+        t_max: 80.0
+        t_min: 0.002
+      tokenizer:
+        _target_: <class 'cosmos_predict2.tokenizers.tokenizer.TokenizerInterface'>
+        chunk_duration: '81'
+        load_mean_std: 'False'
+        name: tokenizer
+        temporal_window: '16'
+        vae_pth: /home/nvidia/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth
+model_parallel:
+  _cpu_offloading_context: null
+  async_tensor_model_parallel_allreduce: false
+  autocast_dtype: torch.float32
+  barrier_with_L1_time: true
+  batch_p2p_comm: true
+  batch_p2p_sync: true
+  bf16: false
+  context_parallel_size: 1
+  cpu_offloading: false
+  cpu_offloading_activations: false
+  cpu_offloading_num_layers: 0
+  cpu_offloading_weights: false
+  cross_entropy_fusion_impl: native
+  cross_entropy_loss_fusion: false
+  deallocate_pipeline_outputs: false
+  defer_embedding_wgrad_compute: false
+  deterministic_mode: false
+  enable_autocast: false
+  expert_model_parallel_size: 1
+  expert_tensor_parallel_size: 1
+  finalize_model_grads_func: null
+  fp16: false
+  grad_scale_func: null
+  grad_sync_func: null
+  gradient_accumulation_fusion: false
+  hierarchical_context_parallel_sizes: null
+  microbatch_group_size_per_vp_stage: 1
+  moe_extended_tp: false
+  no_sync_func: null
+  num_microbatches_with_partial_activation_checkpoints: null
+  overlap_p2p_comm: false
+  overlap_p2p_comm_warmup_flush: false
+  param_sync_func: null
+  params_dtype: torch.float32
+  perform_initialization: true
+  pipeline_dtype: null
+  pipeline_model_parallel_comm_backend: null
+  pipeline_model_parallel_size: 1
+  pipeline_model_parallel_split_rank: null
+  sequence_parallel: false
+  tensor_model_parallel_size: 1
+  timers: null
+  tp_comm_atomic_ag: false
+  tp_comm_atomic_rs: false
+  tp_comm_bootstrap_backend: nccl
+  tp_comm_bulk_dgrad: true
+  tp_comm_bulk_wgrad: true
+  tp_comm_overlap: false
+  tp_comm_overlap_ag: true
+  tp_comm_overlap_disable_fc1: false
+  tp_comm_overlap_disable_qkv: false
+  tp_comm_overlap_rs: true
+  tp_comm_overlap_rs_dgrad: false
+  tp_comm_split_ag: true
+  tp_comm_split_rs: true
+  use_cpu_initialization: false
+  use_ring_exchange_p2p: false
+  use_te_rng_tracker: false
+  variable_seq_lengths: false
+  virtual_pipeline_model_parallel_size: null
+  wgrad_deferral_limit: 0
+optimizer:
+  _target_: <function get_base_optimizer at 0x77eaa4382710>
+  betas:
+  - '0.9'
+  - '0.99'
+  capturable: 'True'
+  eps: 1e-08
+  lr: '9.999999999999999e-05'
+  master_weights: 'True'
+  model: null
+  optim_type: fusedadam
+  weight_decay: '0.1'
+scheduler:
+  _target_: <class 'cosmos_predict2.functional.lr_scheduler.LambdaLinearScheduler'>
+  cycle_lengths:
+  - '500000'
+  f_max:
+  - '1'
+  f_min:
+  - '0.2'
+  f_start:
+  - 1e-06
+  verbosity_interval: '0'
+  warm_up_steps:
+  - '1000'
+trainer:
+  callbacks:
+    device_monitor:
+      _target_: <class 'cosmos_predict2.callbacks.device_monitor.DeviceMonitor'>
+      every_n: '1000'
+      log_memory_detail: 'True'
+      step_size: '1'
+    ema:
+      _target_: <class 'imaginaire.utils.callback.EMAModelCallback'>
+      config: null
+      trainer: null
+    grad_clip:
+      _target_: <class 'cosmos_predict2.callbacks.grad_clip.GradClip'>
+      clip_norm: '10.0'
+      force_finite: 'True'
+      log_wandb: 'False'
+    iter_speed:
+      _target_: <class 'cosmos_predict2.callbacks.iter_speed.IterSpeed'>
+      every_n: '1000'
+      hit_thres: '5'
+    low_prec:
+      _target_: <class 'imaginaire.utils.callback.LowPrecisionCallback'>
+      config: null
+      trainer: null
+      update_iter: '1'
+    manual_gc:
+      _target_: <class 'imaginaire.callbacks.manual_gc.ManualGarbageCollection'>
+      every_n: '5'
+      warm_up: '5'
+    progress_bar:
+      _target_: <class 'imaginaire.utils.callback.ProgressBarCallback'>
+      config: null
+      trainer: null
+    wandb:
+      _target_: <class 'cosmos_predict2.callbacks.wandb_logger.WandBLogger'>
+      mode: null
+  cudnn:
+    benchmark: 'True'
+    deterministic: 'False'
+  ddp:
+    broadcast_buffers: 'True'
+    find_unused_parameters: 'False'
+    static_graph: 'True'
+  distributed_parallelism: ddp
+  grad_accum_iter: '8'
+  grad_scaler_args:
+    enabled: 'False'
+  logging_iter: '1000'
+  max_iter: '8000'
+  max_val_iter: '16'
+  memory_format: torch.preserve_format
+  profiling:
+    enable_memory_snapshot: 'False'
+    enable_profiling: 'False'
+    first_n_rank: '4'
+    profile_freq: '1'
+    profile_memory: 'True'
+    record_shape: 'True'
+    with_modules: 'True'
+    with_stack: 'True'
+  run_validation: 'True'
+  seed: '0'
+  timeout_period: '999999999'
+  type: <class 'imaginaire.trainer.ImaginaireTrainer'>
+  validation_iter: '100'
+video_dataset_train: null
+video_dataset_val: null
+world2action_pipe:
+  _target_: cosmos_predict2.configs.config_world2action.World2ActionPipelineConfig
+  ema:
+    enabled: false
+    iteration_shift: 0
+    rate: 0.1
+  net:
+    _target_: <class 'cosmos_predict2.models.world2action_dit.World2ActionDIT'>
+    adaln_lora_dim: '128'
+    atten_backend: flash_attn_no_cp
+    crossattn_emb_channels: '2048'
+    in_channels: '6'
+    max_horizon: '91'
+    mlp_ratio: '4.0'
+    model_channels: '512'
+    num_blocks: '12'
+    num_heads: '8'
+    out_channels: '6'
+    pair_timestep_feature_rank: '512'
+    sac_config:
+      every_n_blocks: 1
+      mode: none
+    use_adaln_lora: 'True'
+  precision: bfloat16
+  scheduler:
+    alpha: 1.0
+    beta: 1.0
+    num_denoising_steps: 10

action_decoder/chain/w2a_so101_task2_combined_lora32_iter2500_75m_2gpu_bsz2_accum8_val5/stdout.log

@@ -0,0 +1,102 @@
+[05-17 22:03:26|INFO|imaginaire/trainer.py:116:__init__] Config:
+* model: {'config': {'validation_mode': 'loss_only', 'train_architecture': 'base', 'lora_rank': 16, 'lora_alpha': 16, 'lora_target_modules': 'q_proj,k_proj,v_proj,output_proj,mlp.layer1,mlp.layer2', 'init_lora_weights': True, 'precision': 'bfloat16', 'loss_reduce': 'mean', 'loss_scale': 10.0, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, 'action_dit_path': '', 'video_dit_path': '/home/nvidia/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt', 'pipe_config': {'precision': 'bfloat16', 'scheduler': {'alpha': 1.0, 'beta': 1.0, 'num_denoising_steps': 10}, 'net': {'max_horizon': 91, 'in_channels': 6, 'out_channels': 6, 'model_channels': 512, 'num_blocks': 12, 'num_heads': 8, 'mlp_ratio': 4.0, 'atten_backend': 'flash_attn_no_cp', 'crossattn_emb_channels': 2048, 'use_adaln_lora': True, 'adaln_lora_dim': 128, 'pair_timestep_feature_rank': 512, 'sac_config': {'mode': 'none', 'every_n_blocks': 1}, '_target_': <class 'cosmos_predict2.models.world2action_dit.World2ActionDIT'>}, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, '_target_': 'cosmos_predict2.configs.config_world2action.World2ActionPipelineConfig', 'xattn_layer_idx': 20}, 'video_pipe_config': {'adjust_video_noise': True, 'conditioner': {'fps': {'output_key': 'fps', 'dropout_rate': 0.0, 'dtype': None, 'input_key': 'fps', '_target_': <class 'cosmos_predict2.conditioner.ReMapkey'>}, 'padding_mask': {'output_key': 'padding_mask', 'dropout_rate': 0.0, 'dtype': None, 'input_key': 'padding_mask', '_target_': <class 'cosmos_predict2.conditioner.ReMapkey'>}, 'text': {'dropout_rate': 0.0, 'input_key': ['obs/language_embedding'], '_target_': <class 'cosmos_predict2.conditioner.TextAttr'>}, 'use_video_condition': {'output_key': 'use_video_condition', 'dropout_rate': 0.0, 'input_key': 'fps', '_target_': <class 'cosmos_predict2.conditioner.BooleanFlag'>}, '_target_': <class 'cosmos_predict2.conditioner.VideoConditioner'>}, 'conditioning_strategy': 'frame_replace', 'min_num_conditional_frames': 1, 'max_num_conditional_frames': 2, 'sigma_conditional': 0.0001, 'net': {'max_img_h': 240, 'max_img_w': 240, 'max_frames': 128, 'in_channels': 16, 'out_channels': 16, 'patch_spatial': 2, 'patch_temporal': 1, 'concat_padding_mask': True, 'model_channels': 2048, 'num_blocks': 28, 'num_heads': 16, 'atten_backend': 'minimal_a2a', 'pos_emb_cls': 'rope3d', 'pos_emb_learnable': True, 'pos_emb_interpolation': 'crop', 'use_adaln_lora': True, 'adaln_lora_dim': 256, 'rope_h_extrapolation_ratio': 3.0, 'rope_w_extrapolation_ratio': 3.0, 'rope_t_extrapolation_ratio': 1.0, 'extra_per_block_abs_pos_emb': False, 'rope_enable_fps_modulation': False, 'sac_config': {'mode': 'predict2_2b_720', 'every_n_blocks': 1, '_target_': 'cosmos_predict2.models.text2image_dit.SACConfig'}, '_target_': <class 'cosmos_predict2.models.video2world_dit.MinimalV1LVGDiT'>}, 'tokenizer': {'chunk_duration': 81, 'load_mean_std': False, 'temporal_window': 16, 'name': 'tokenizer', 'vae_pth': '/home/nvidia/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth', '_target_': <class 'cosmos_predict2.tokenizers.tokenizer.TokenizerInterface'>}, 'guardrail_config': {'checkpoint_dir': '/home/nvidia/pdt-mimic/mimic-video/model/checkpoints', 'offload_model_to_cpu': True, 'enabled': False}, 'precision': 'bfloat16', 'rectified_flow_t_scaling_factor': 1.0, 'rectified_flow_loss_weight_uniform': True, 'resize_online': False, 'resolution': '480', 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0, '_target_': 'cosmos_predict2.configs.defaults.ema.EMAConfig'}, 'sigma_data': 1.0, 'state_ch': 16, 'state_t': 16, 'text_encoder': {'cls': <TextEncoderClass.T5: 't5'>, 't5': {'ckpt_path': '/home/nvidia/pdt-mimic/mimic-video/model/checkpoints/text_encoder/t5-11b', 'num_tokens': 512, 'embed_dim': 1024}}, 'input_video_key': 'video', 'input_image_key': 'images', 'timestamps': {'nfe': 35, 't_min': 0.002, 't_max': 80.0, 'order': 7.0, 'is_forward': False}}, 'fsdp_shard_size': 0, 'data_config': {'config_name': 'so101_lerobot', '_target_': <function get_data_config at 0x77eab678a680>}, '_target_': 'cosmos_predict2.models.world2action_model.World2ActionModelConfig'}, '_target_': <class 'cosmos_predict2.models.world2action_model.World2ActionModel'>}
+* world2action_pipe: {'precision': 'bfloat16', 'scheduler': {'alpha': 1.0, 'beta': 1.0, 'num_denoising_steps': 10}, 'net': {'max_horizon': 91, 'in_channels': 6, 'out_channels': 6, 'model_channels': 512, 'num_blocks': 12, 'num_heads': 8, 'mlp_ratio': 4.0, 'atten_backend': 'flash_attn_no_cp', 'crossattn_emb_channels': 2048, 'use_adaln_lora': True, 'adaln_lora_dim': 128, 'pair_timestep_feature_rank': 512, 'sac_config': {'mode': 'none', 'every_n_blocks': 1}, '_target_': <class 'cosmos_predict2.models.world2action_dit.World2ActionDIT'>}, 'ema': {'enabled': False, 'rate': 0.1, 'iteration_shift': 0}, '_target_': 'cosmos_predict2.configs.config_world2action.World2ActionPipelineConfig'}
+* optimizer: {'optim_type': 'fusedadam', 'model': None, 'lr': 9.999999999999999e-05, 'weight_decay': 0.1, 'betas': [0.9, 0.99], 'eps': 1e-08, 'master_weights': True, 'capturable': True, '_target_': <function get_base_optimizer at 0x77eaa4382710>}
+* scheduler: {'verbosity_interval': 0, 'warm_up_steps': [1000], 'cycle_lengths': [500000], 'f_start': [1e-06], 'f_max': [1], 'f_min': [0.2], '_target_': <class 'cosmos_predict2.functional.lr_scheduler.LambdaLinearScheduler'>}
+* data_config: {'config_name': 'so101_lerobot', '_target_': <function get_data_config at 0x77eab678a680>}
+* video_dataset_train: None
+* video_dataset_val: None
+* dataloader_train: {'batch_size': 2, 'shuffle': None, 'sampler': {'seed': 0, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': None, 'seed': 42, 'train': True, 'cache_dir': None, 'video_backend': 'pyav', 'repo_id': 'push-that-thing/task_1', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, 'num_replicas': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_world_size at 0x77ed19e6d5a0>}, 'rank': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_rank at 0x77ed19e6d6c0>}, 'shuffle': True, '_target_': <class 'cosmos_predict2.data.resumable_sampler.ResumableDistributedSampler'>}, 'batch_sampler': None, 'num_workers': 8, 'collate_fn': None, 'pin_memory': True, 'drop_last': True, 'timeout': 0, 'worker_init_fn': None, 'multiprocessing_context': None, 'generator': None, 'prefetch_factor': 8, 'persistent_workers': True, 'pin_memory_device': '', 'in_order': True, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': 0.05, 'seed': 42, 'train': True, 'cache_dir': None, 'video_backend': 'pyav', 'repo_id': 'push-that-thing/task_2', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, '_target_': <class 'torch.utils.data.dataloader.DataLoader'>}
+* dataloader_val: {'batch_size': 1, 'shuffle': None, 'sampler': {'seed': 0, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': None, 'seed': 42, 'train': False, 'cache_dir': None, 'video_backend': 'pyav', 'repo_id': 'push-that-thing/task_1', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, 'num_replicas': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_world_size at 0x77ed19e6d5a0>}, 'rank': {'with_context_parallel': False, 'partial_data_parallel': False, '_target_': <function get_data_parallel_rank at 0x77ed19e6d6c0>}, 'shuffle': False, '_target_': <class 'cosmos_predict2.data.resumable_sampler.ResumableDistributedSampler'>}, 'batch_sampler': None, 'num_workers': 0, 'collate_fn': None, 'pin_memory': False, 'drop_last': False, 'timeout': 0, 'worker_init_fn': None, 'multiprocessing_context': None, 'generator': None, 'prefetch_factor': None, 'persistent_workers': False, 'pin_memory_device': '', 'in_order': True, 'dataset': {'root': None, 'image_key': 'observation.images.front', 'state_key': 'observation.state', 'action_key': 'action', 'obs_image_horizon': 5, 'action_image_horizon': 56, 'action_lowdim_horizon': 90, 'target_fps': 5.0, 'lowdim_target_fps': 30.0, 'action_shift_s': 0.2, 'image_resize': [480, 640], 'language_embeddings_path': None, 't5_encoder_ckpt': None, 'num_val_episodes': 1, 'val_ratio': 0.05, 'seed': 42, 'train': False, 'cache_dir': None, 'video_backend': 'pyav', 'repo_id': 'push-that-thing/task_2', '_target_': <class 'cosmos_predict2.data.action.dataset_lerobot.LeRobotDataset'>}, '_target_': <class 'torch.utils.data.dataloader.DataLoader'>}
+* job:
+   * project: vam
+   * group: so101
+   * name: w2a_so101_task2_combined_lora32_iter2500_75m_2gpu_bsz2_accum8_val5
+* trainer:
+   * type: <class 'imaginaire.trainer.ImaginaireTrainer'>
+   * callbacks: {'ema': {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.EMAModelCallback'>}, 'progress_bar': {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.ProgressBarCallback'>}, 'low_prec': {'config': None, 'trainer': None, 'update_iter': 1, '_target_': <class 'imaginaire.utils.callback.LowPrecisionCallback'>}, 'iter_speed': {'hit_thres': 5, 'every_n': 1000, '_target_': <class 'cosmos_predict2.callbacks.iter_speed.IterSpeed'>}, 'device_monitor': {'every_n': 1000, 'step_size': 1, 'log_memory_detail': True, '_target_': <class 'cosmos_predict2.callbacks.device_monitor.DeviceMonitor'>}, 'manual_gc': {'warm_up': 5, 'every_n': 5, '_target_': <class 'imaginaire.callbacks.manual_gc.ManualGarbageCollection'>}, 'grad_clip': {'clip_norm': 10.0, 'force_finite': True, 'log_wandb': False, '_target_': <class 'cosmos_predict2.callbacks.grad_clip.GradClip'>}, 'wandb': {'mode': None, '_target_': <class 'cosmos_predict2.callbacks.wandb_logger.WandBLogger'>}}
+   * distributed_parallelism: ddp
+   * ddp:
+      * find_unused_parameters: False
+      * static_graph: True
+      * broadcast_buffers: True
+   * cudnn:
+      * deterministic: False
+      * benchmark: True
+   * seed: 0
+   * grad_scaler_args: {'enabled': False}
+   * max_iter: 8000
+   * max_val_iter: 16
+   * logging_iter: 1000
+   * run_validation: True
+   * validation_iter: 100
+   * timeout_period: 999999999
+   * memory_format: torch.preserve_format
+   * grad_accum_iter: 8
+   * profiling:
+      * enable_profiling: False
+      * enable_memory_snapshot: False
+      * profile_freq: 1
+      * first_n_rank: 4
+      * record_shape: True
+      * profile_memory: True
+      * with_stack: True
+      * with_modules: True
+* model_parallel: ModelParallelConfig(tensor_model_parallel_size=1, pipeline_model_parallel_comm_backend=None, pipeline_model_parallel_size=1, virtual_pipeline_model_parallel_size=None, sequence_parallel=False, context_parallel_size=1, hierarchical_context_parallel_sizes=None, expert_model_parallel_size=1, expert_tensor_parallel_size=1, moe_extended_tp=False, perform_initialization=True, use_cpu_initialization=False, fp16=False, bf16=False, params_dtype=torch.float32, timers=None, finalize_model_grads_func=None, grad_scale_func=None, no_sync_func=None, grad_sync_func=None, param_sync_func=None, deterministic_mode=False, enable_autocast=False, autocast_dtype=torch.float32, num_microbatches_with_partial_activation_checkpoints=None, gradient_accumulation_fusion=False, async_tensor_model_parallel_allreduce=False, use_te_rng_tracker=False, tp_comm_overlap=False, tp_comm_bulk_wgrad=True, tp_comm_bulk_dgrad=True, tp_comm_overlap_ag=True, tp_comm_overlap_rs=True, tp_comm_overlap_rs_dgrad=False, tp_comm_split_ag=True, tp_comm_atomic_ag=False, tp_comm_split_rs=True, tp_comm_atomic_rs=False, cross_entropy_loss_fusion=False, cross_entropy_fusion_impl='native', tp_comm_overlap_disable_qkv=False, tp_comm_overlap_disable_fc1=False, tp_comm_bootstrap_backend='nccl', pipeline_dtype=None, variable_seq_lengths=False, overlap_p2p_comm=False, batch_p2p_comm=True, batch_p2p_sync=True, use_ring_exchange_p2p=False, deallocate_pipeline_outputs=False, defer_embedding_wgrad_compute=False, wgrad_deferral_limit=0, pipeline_model_parallel_split_rank=None, overlap_p2p_comm_warmup_flush=False, microbatch_group_size_per_vp_stage=1, cpu_offloading=False, cpu_offloading_num_layers=0, _cpu_offloading_context=None, cpu_offloading_activations=False, cpu_offloading_weights=False, barrier_with_L1_time=True)
+* checkpoint:
+   * type: {'callbacks': None, '_target_': <class 'cosmos_predict2.checkpointer.Checkpointer'>}
+   * dcp_async_mode_enabled: False
+   * save_iter: 100
+   * load_path: 
+   * load_training_state: False
+   * only_load_scheduler_state: False
+   * strict_resume: True
+   * jit:
+      * enabled: False
+      * input_shape: None
+      * device: cuda
+      * dtype: bfloat16
+      * strict: True
+   * verbose: True
+   * keys_not_to_resume: []
+   * broadcast_via_filesystem: False
+   * load_ema_to_reg: False
+   * dcp_allow_mismatched_size: False
+* defaults: ['_self_', {'data_config': None}, {'video_dataset_train': None}, {'video_dataset_val': None}, {'dataloader_train': None}, {'dataloader_val': None}, {'world2action_pipe': None}, {'optimizer': 'fusedadamw'}, {'scheduler': 'constant'}, {'model': None}, {'callbacks': ['basic']}, {'net': None}, {'ema': None}, {'checkpoint': None}, {'ckpt_type': None}, {'experiment': None}]
+[05-17 22:03:26|WARNING|imaginaire/utils/misc.py:127:print_environ_variables] Environment variable TORCH_HOME not set!
+[05-17 22:03:26|WARNING|imaginaire/utils/misc.py:127:print_environ_variables] Environment variable IMAGINAIRE_OUTPUT_ROOT not set!
+[05-17 22:03:26|INFO|imaginaire/utils/misc.py:139:set_random_seed] Using random seed 0.
+[05-17 22:03:26|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback ema: {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.EMAModelCallback'>}
+[05-17 22:03:26|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback progress_bar: {'config': None, 'trainer': None, '_target_': <class 'imaginaire.utils.callback.ProgressBarCallback'>}
+[05-17 22:03:26|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback low_prec: {'config': None, 'trainer': None, 'update_iter': 1, '_target_': <class 'imaginaire.utils.callback.LowPrecisionCallback'>}
+[05-17 22:03:26|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback iter_speed: {'hit_thres': 5, 'every_n': 1000, '_target_': <class 'cosmos_predict2.callbacks.iter_speed.IterSpeed'>}
+[05-17 22:03:26|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback device_monitor: {'every_n': 1000, 'step_size': 1, 'log_memory_detail': True, '_target_': <class 'cosmos_predict2.callbacks.device_monitor.DeviceMonitor'>}
+[05-17 22:03:26|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback manual_gc: {'warm_up': 5, 'every_n': 5, '_target_': <class 'imaginaire.callbacks.manual_gc.ManualGarbageCollection'>}
+[05-17 22:03:26|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback grad_clip: {'clip_norm': 10.0, 'force_finite': True, 'log_wandb': False, '_target_': <class 'cosmos_predict2.callbacks.grad_clip.GradClip'>}
+[05-17 22:03:26|CRITICAL|imaginaire/utils/callback.py:77:__init__] Instantiating callback wandb: {'mode': None, '_target_': <class 'cosmos_predict2.callbacks.wandb_logger.WandBLogger'>}
+[05-17 22:03:28|CRITICAL|cosmos_predict2/models/world2action_model.py:106:__init__] Using mean loss reduce with loss scale 10.0
+[05-17 22:03:28|WARNING|cosmos_predict2/pipelines/world2action.py:71:from_config] precision torch.bfloat16
+[05-17 22:03:28|WARNING|cosmos_predict2/pipelines/world2action.py:101:from_config] dit_path not provided, initializing DiT with random weights
+[05-17 22:03:28|WARNING|cosmos_predict2/pipelines/video2world.py:277:from_config] precision torch.bfloat16
+[05-17 22:03:29|INFO|cosmos_predict2/tokenizers/tokenizer.py:687:_video_vae] Loading /home/nvidia/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth
+[05-17 22:03:29|SUCCESS|cosmos_predict2/tokenizers/tokenizer.py:689:_video_vae] Successfully loaded /home/nvidia/pdt-mimic/mimic-video/model/checkpoints/video_backbone//tokenizer/tokenizer.pth
+[05-17 22:03:29|INFO|imaginaire/utils/distributed.py:430:sync_model_states] Synchronizing model states from rank 0 to all ranks in process group [0, 1].
+[05-17 22:03:29|INFO|cosmos_predict2/pipelines/video2world.py:335:from_config] Loading DiT from /home/nvidia/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt
+[05-17 22:03:36|SUCCESS|cosmos_predict2/pipelines/video2world.py:354:from_config] Successfully loaded DiT from /home/nvidia/pdt-mimic/mimic-video/model/checkpoints/video_backbone/v2w_push_combined_lora_rank32_lr1.778e-04_bsz32_iter_000002500_fused.pt
+[05-17 22:03:37|INFO|cosmos_predict2/models/world2action_model.py:148:__init__] Total parameters: 2.04B, Frozen parameters: 1,956,413,440, Trainable parameters: 80,859,000
+[05-17 22:03:37|INFO|cosmos_predict2/models/world2action_model.py:167:__init__] FSDP (Fully Sharded Data Parallel) is disabled.
+[05-17 22:03:38|CRITICAL|cosmos_predict2/utils/optim_instantiate_dtensor.py:49:get_base_optimizer] total num parameters : 80,859,000
+[05-17 22:03:38|WARNING|cosmos_predict2/utils/fused_adam_dtensor.py:103:__init__] FusedAdam master_weights: True capturable: True
+[05-17 22:03:38|INFO|cosmos_predict2/checkpointer.py:288:load] Training from scratch.
+[05-17 22:03:38|INFO|imaginaire/utils/misc.py:266:wrapper] Time spent on checkpoint loading: 0.0003 seconds
+[05-17 22:03:38|CRITICAL|imaginaire/trainer.py:179:train] Distributed parallelism mode: ddp
+[05-17 22:03:39|INFO|imaginaire/trainer.py:187:train] Starting training...
+[05-17 22:03:39|INFO|cosmos_predict2/callbacks/device_monitor.py:92:on_train_start] DeviceMonitor callback: local_dir: checkpoints/vam/so101/w2a_so101_task2_combined_lora32_iter2500_75m_2gpu_bsz2_accum8_val5/DeviceMonitor
+[05-17 22:03:40|INFO|imaginaire/trainer.py:388:validate] Validating at iteration 0...
+[05-17 22:04:30|INFO|imaginaire/trainer.py:193:train] Initial validation done.
+[05-17 22:04:35|WARNING|imaginaire/utils/distributed.py:284:ddp_sync_grad] DDP static_graph=True is incompatible with sync_grad(). Performance will be reduced.
+[05-17 22:05:04|INFO|cosmos_predict2/callbacks/iter_speed.py:51:on_training_step_end] Iteration 1: Hit counter: 1/5 | Loss: 31.3163 | Time: 98.58s
+[05-17 22:05:22|INFO|cosmos_predict2/callbacks/iter_speed.py:51:on_training_step_end] Iteration 2: Hit counter: 2/5 | Loss: 23.0725 | Time: 17.80s
+[05-17 22:05:40|INFO|cosmos_predict2/callbacks/iter_speed.py:51:on_training_step_end] Iteration 3: Hit counter: 3/5 | Loss: 49.5040 | Time: 17.61s
+[05-17 22:05:57|INFO|cosmos_predict2/callbacks/iter_speed.py:51:on_training_step_end] Iteration 4: Hit counter: 4/5 | Loss: 31.5268 | Time: 17.69s
+[05-17 22:06:15|INFO|cosmos_predict2/callbacks/iter_speed.py:51:on_training_step_end] Iteration 5: Hit counter: 5/5 | Loss: 26.1954 | Time: 17.56s