Use current module names in model configs and checkpoint metadata; link the paper

README.md

@@ -11,7 +11,7 @@ tags:
 
 Checkpoints for **PrismWAM: Decoupling State Representations in World Action Models**.
 
-[Code](https://github.com/Prism-WAM/PrismWAM) · [Project page](https://prism-wam.github.io/)
+[Code](https://github.com/Prism-WAM/PrismWAM) · [Project page](https://prism-wam.github.io/) · [Paper](https://prism-wam.github.io/static/paper/PrismWAM.pdf)
 
 | Folder | Benchmark | Actions | Result |
 | --- | --- | --- | ---: |

libero/config.yaml

@@ -1,15 +1,11 @@
 model:
   _target_: prismwam.models.wan22.rgb_depth_action_depth_parent.create_prismwam_rgb_depth_action
   model_id: Wan-AI/Wan2.2-TI2V-5B
-  rgb_model_id: null
   tokenizer_model_id: Wan-AI/Wan2.1-T2V-1.3B
   tokenizer_max_len: 128
   load_text_encoder: false
   redirect_common_files: true
   architecture_name: PrismWAM
-  rgb_vae_family: wan22_vae48
-  rgb_dit_io_adaptation: null
-  rgb_attention_adapter: auto
   rgb_input_size:
   - 224
   - 448
@@ -18,15 +14,7 @@ model:
   rgb_fast: true
   rgb_history_num_frames: 5
   rgb_history_frame_stride: 4
-  depth_history_num_frames: 1
-  depth_history_frame_stride: 4
-  rgb_tower_dropout_prob: 0.0
-  rgb_future_dropout_prob: 0.0
   share_rgb_depth_blocks: true
-  rgb_depth_shared_core_layer_range: null
-  rgb_temporal_rope_positions: null
-  rgb_to_depth_frame_attention: all_rgb_frames
-  training_sigma_sampling: independent_per_tower
   video_use_text_embedding: true
   proprio_dim: 8
   mot_checkpoint_mixed_attn: true
@@ -34,15 +22,12 @@ model:
   skip_rgb_dit_load_from_pretrain: true
   skip_action_dit_load_from_pretrain: true
   checkpoint_init_only: true
-  checkpoint_action_dim_transition: strict
-  checkpoint_action_dim_transition_parent: null
   depth_vae_require_latent_stats: true
   depth_dit_checkpoint_path: null
   depth_vae_checkpoint_path: depth_vae.pt
   rgb_vae_checkpoint_path: rgb_vae.safetensors
   action_dit_pretrained_path: null
   depth_dit_config:
-    has_image_input: false
     patch_size:
     - 1
     - 2
@@ -58,16 +43,11 @@ model:
     num_layers: 30
     eps: 1.0e-06
     seperated_timestep: true
-    require_clip_embedding: false
     require_vae_embedding: false
     fuse_vae_embedding_in_latents: true
     use_gradient_checkpointing: true
     video_attention_mask_mode: first_frame_causal
-    action_conditioned: false
-    action_dim: 7
-    action_group_causal_mask_mode: group_diagonal
   rgb_dit_config:
-    has_image_input: false
     patch_size:
     - 1
     - 2
@@ -83,14 +63,10 @@ model:
     num_layers: 30
     eps: 1.0e-06
     seperated_timestep: true
-    require_clip_embedding: false
     require_vae_embedding: false
     fuse_vae_embedding_in_latents: true
     use_gradient_checkpointing: true
     video_attention_mask_mode: first_frame_causal
-    action_conditioned: false
-    action_dim: 7
-    action_group_causal_mask_mode: group_diagonal
   action_dit_config:
     action_dim: 7
     hidden_dim: 1024
@@ -181,14 +157,8 @@ data:
         - true
         - true
         - false
-      action_state_transforms: null
-      use_stepwise_action_norm: false
-      norm_default_mode: min/max
-      norm_exception_mode: null
       action_state_merger:
         _target_: prismwam.datasets.lerobot.transforms.action_state_merger.ConcatLeftAlign
-      train_transforms: []
-      val_transforms: []
 deployment:
   action_horizon: 32
   replan_steps: 16
@@ -199,4 +169,3 @@ deployment:
   benchmark: libero
   portable_bundle: true
   norm_stats_path: norm_stats.json
-eval_num_inference_steps: 10

libero/policy.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9972469886e613b944b0aa3273f2b5df3fe0b5e8bb8641b4a4956426b1bc537e
-size 12222029803
+oid sha256:d88b7bd35538d6c325ac66c24f00bdb62a5ae0e8227b5aa190065e444134cdb9
+size 12222027990

robodojo/config.yaml

@@ -1,15 +1,11 @@
 model:
-  _target_: prismwam.models.wan22.rgb_memory_duet.create_prismwam_rgb_depth_action
+  _target_: prismwam.models.wan22.rgb_memory.create_prismwam_rgb_depth_action
   model_id: Wan-AI/Wan2.2-TI2V-5B
-  rgb_model_id: null
   tokenizer_model_id: Wan-AI/Wan2.1-T2V-1.3B
   tokenizer_max_len: 128
   load_text_encoder: false
   redirect_common_files: true
   architecture_name: PrismWAM
-  rgb_vae_family: wan22_vae48
-  rgb_dit_io_adaptation: null
-  rgb_attention_adapter: auto
   rgb_input_size:
   - 384
   - 320
@@ -18,15 +14,7 @@ model:
   rgb_fast: true
   rgb_history_num_frames: 3
   rgb_history_frame_stride: 4
-  depth_history_num_frames: 1
-  depth_history_frame_stride: 4
-  rgb_tower_dropout_prob: 0.0
-  rgb_future_dropout_prob: 0.0
   share_rgb_depth_blocks: true
-  rgb_depth_shared_core_layer_range: null
-  rgb_temporal_rope_positions: null
-  rgb_to_depth_frame_attention: all_rgb_frames
-  training_sigma_sampling: independent_per_tower
   video_use_text_embedding: true
   proprio_dim: 14
   mot_checkpoint_mixed_attn: true
@@ -34,15 +22,12 @@ model:
   skip_rgb_dit_load_from_pretrain: true
   skip_action_dit_load_from_pretrain: true
   checkpoint_init_only: true
-  checkpoint_action_dim_transition: strict
-  checkpoint_action_dim_transition_parent: null
   depth_vae_require_latent_stats: true
   depth_dit_checkpoint_path: null
   depth_vae_checkpoint_path: depth_vae.pt
   rgb_vae_checkpoint_path: rgb_vae.safetensors
   action_dit_pretrained_path: null
   depth_dit_config:
-    has_image_input: false
     patch_size:
     - 1
     - 2
@@ -58,16 +43,11 @@ model:
     num_layers: 30
     eps: 1.0e-06
     seperated_timestep: true
-    require_clip_embedding: false
     require_vae_embedding: false
     fuse_vae_embedding_in_latents: true
     use_gradient_checkpointing: true
     video_attention_mask_mode: first_frame_causal
-    action_conditioned: false
-    action_dim: 14
-    action_group_causal_mask_mode: group_diagonal
   rgb_dit_config:
-    has_image_input: false
     patch_size:
     - 1
     - 2
@@ -83,14 +63,10 @@ model:
     num_layers: 30
     eps: 1.0e-06
     seperated_timestep: true
-    require_clip_embedding: false
     require_vae_embedding: false
     fuse_vae_embedding_in_latents: true
     use_gradient_checkpointing: true
     video_attention_mask_mode: first_frame_causal
-    action_conditioned: false
-    action_dim: 14
-    action_group_causal_mask_mode: group_diagonal
   action_dit_config:
     action_dim: 14
     hidden_dim: 1024
@@ -130,8 +106,6 @@ model:
     sparse_recent_action_offset: 32
     sparse_frame_encoding: independent
     rgb_proprio_context: false
-    training_cache_mode: full
-    checkpoint_recompute: false
     conditioning_latent_noise_prob: 0.2
     conditioning_latent_noise_stream: rgb
     rgb_conditioning_latent_noise_prob: 0.2
@@ -145,8 +119,6 @@ data:
     frame_stride: 4
     rgb_history_num_frames: 3
     rgb_history_frame_stride: 4
-    depth_history_num_frames: 1
-    depth_history_frame_stride: 4
     cameras:
     - cam_head
     - cam_left_wrist
@@ -171,4 +143,3 @@ deployment:
   benchmark: robodojo
   portable_bundle: true
   statistics_path: action_stats.pt
-eval_num_inference_steps: 10

robodojo/policy.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:39e8fa611c9adc14df6a9c97cc4f06742bb46a6acd85146ea4c89d09adb087e3
-size 12222108523
+oid sha256:ddbbc70515a03d876740110e3bb6112a50bbd8724dceeabc085dd28ab04cdbe0
+size 12222106646

robotwin/config.yaml

@@ -1,15 +1,11 @@
 model:
   _target_: prismwam.models.wan22.rgb_depth_action_depth_parent.create_prismwam_rgb_depth_action
   model_id: Wan-AI/Wan2.2-TI2V-5B
-  rgb_model_id: null
   tokenizer_model_id: Wan-AI/Wan2.1-T2V-1.3B
   tokenizer_max_len: 128
   load_text_encoder: false
   redirect_common_files: true
   architecture_name: PrismWAM
-  rgb_vae_family: wan22_vae48
-  rgb_dit_io_adaptation: null
-  rgb_attention_adapter: auto
   rgb_input_size:
   - 384
   - 320
@@ -18,15 +14,7 @@ model:
   rgb_fast: true
   rgb_history_num_frames: 5
   rgb_history_frame_stride: 4
-  depth_history_num_frames: 1
-  depth_history_frame_stride: 4
-  rgb_tower_dropout_prob: 0.0
-  rgb_future_dropout_prob: 0.0
   share_rgb_depth_blocks: true
-  rgb_depth_shared_core_layer_range: null
-  rgb_temporal_rope_positions: null
-  rgb_to_depth_frame_attention: all_rgb_frames
-  training_sigma_sampling: independent_per_tower
   video_use_text_embedding: true
   proprio_dim: 14
   mot_checkpoint_mixed_attn: true
@@ -34,15 +22,12 @@ model:
   skip_rgb_dit_load_from_pretrain: true
   skip_action_dit_load_from_pretrain: true
   checkpoint_init_only: true
-  checkpoint_action_dim_transition: strict
-  checkpoint_action_dim_transition_parent: null
   depth_vae_require_latent_stats: true
   depth_dit_checkpoint_path: null
   depth_vae_checkpoint_path: depth_vae.pt
   rgb_vae_checkpoint_path: rgb_vae.safetensors
   action_dit_pretrained_path: null
   depth_dit_config:
-    has_image_input: false
     patch_size:
     - 1
     - 2
@@ -58,16 +43,11 @@ model:
     num_layers: 30
     eps: 1.0e-06
     seperated_timestep: true
-    require_clip_embedding: false
     require_vae_embedding: false
     fuse_vae_embedding_in_latents: true
     use_gradient_checkpointing: true
     video_attention_mask_mode: first_frame_causal
-    action_conditioned: false
-    action_dim: 14
-    action_group_causal_mask_mode: group_diagonal
   rgb_dit_config:
-    has_image_input: false
     patch_size:
     - 1
     - 2
@@ -83,14 +63,10 @@ model:
     num_layers: 30
     eps: 1.0e-06
     seperated_timestep: true
-    require_clip_embedding: false
     require_vae_embedding: false
     fuse_vae_embedding_in_latents: true
     use_gradient_checkpointing: true
     video_attention_mask_mode: first_frame_causal
-    action_conditioned: false
-    action_dim: 14
-    action_group_causal_mask_mode: group_diagonal
   action_dit_config:
     action_dim: 14
     hidden_dim: 1024
@@ -130,8 +106,6 @@ data:
     frame_stride: 4
     rgb_history_num_frames: 5
     rgb_history_frame_stride: 4
-    depth_history_num_frames: 1
-    depth_history_frame_stride: 4
     cameras:
     - head_camera
     - left_camera
@@ -159,4 +133,3 @@ deployment:
   benchmark: robotwin
   portable_bundle: true
   statistics_path: action_stats.pt
-eval_num_inference_steps: 10

robotwin/policy.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f6daab3a380a9fc203779c4fdaaa9a757eba6dcfc28e5289dbb51c47f4f64ed1
-size 12222107755
+oid sha256:c6ddd4ed984f69544ccd4ea1c9eadf171a709b1019cf51576c5b9f9af1581143
+size 12222105942