Hugging Face's logo

AAyano
/
gate_setting2_chunksize25_batch32_from20000

Robotics
Transformers
Safetensors
openvla
feature-extraction
vla
openvla-oft
xarm
task-conditioned-gate
custom_code
Model card Files
xet
 Community
AAyano commited on
Commit
f4bdec1
·
verified ·
1 Parent(s): e049115

Add text-token-gate OpenVLA-OFT setting2 (10k fresh on 20k OFT backbone)

Browse files
Files changed (26) hide show
  1. README.md +35 -0
  2. action_head--10000_checkpoint.pt +3 -0
  3. added_tokens.json +3 -0
  4. config.json +3192 -0
  5. configuration_prismatic.py +319 -0
  6. dataset_statistics.json +121 -0
  7. generation_config.json +7 -0
  8. lora_adapter/README.md +202 -0
  9. lora_adapter/adapter_config.json +44 -0
  10. lora_adapter/adapter_model.safetensors +3 -0
  11. model-00001-of-00004.safetensors +3 -0
  12. model-00002-of-00004.safetensors +3 -0
  13. model-00003-of-00004.safetensors +3 -0
  14. model-00004-of-00004.safetensors +3 -0
  15. model.safetensors.index.json +1011 -0
  16. modeling_prismatic.py +1788 -0
  17. oft_training_config.json +15 -0
  18. preprocessor_config.json +114 -0
  19. processing_prismatic.py +257 -0
  20. processor_config.json +6 -0
  21. proprio_projector--10000_checkpoint.pt +3 -0
  22. special_tokens_map.json +30 -0
  23. tokenizer.json +0 -0
  24. tokenizer.model +3 -0
  25. tokenizer_config.json +53 -0
  26. vision_backbone--10000_checkpoint.pt +3 -0
README.md ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ base_model: AAyano/oft_setting2_chunksize25_batch32_20k
4
+ tags:
5
+ - robotics
6
+ - vla
7
+ - openvla
8
+ - openvla-oft
9
+ - xarm
10
+ - task-conditioned-gate
11
+ library_name: transformers
12
+ ---
13
+
14
+ # gate_setting2_chunksize25_batch32_from20000
15
+
16
+ Task-conditioned **text-token-gate** OpenVLA-OFT model for real-world XArm, **setting 2: cup stacking**.
17
+ Trained fresh for **10000 steps** on top of the [oft_setting2_chunksize25_batch32_20k](https://huggingface.co/AAyano/oft_setting2_chunksize25_batch32_20k) OFT backbone (action head / proprio projector / FiLM re-trained; cosine LR 5e-4 with 1000 warmup, effective batch 32, seed 6).
18
+
19
+ ## Gate configuration (baked into `config.json`)
20
+
21
+ - Gazing: `self_attention` (dim 512, 8 heads, 1 layer), gate MLP hidden 512 / depth 1
22
+ - Text pooling: `contrastive_alignment_score` (visual τ 0.1, text τ 0.05), instruction-only pooling, stopword filtering
23
+ - Layer gate: `threshold` mode, threshold 0.3, strength 0.5
24
+ - Budget loss: target 0.4, weight 0.01
25
+
26
+ ## ⚠️ Serving requires the Task-conditioned_Gate codebase
27
+
28
+ The gate modules live in this repo's `modeling_prismatic.py` (loaded via `trust_remote_code`), but the openvla-oft serve flow syncs the local repo's modeling files INTO the checkpoint before loading. Serve with `REPO_DIR` pointing at the [Task-conditioned_Gate](https://github.com/Kaixi66/Task-conditioned_Gate) repo — serving with vanilla openvla-oft silently drops the gate weights.
29
+
30
+ ## Contents
31
+
32
+ - Merged full model (`model-*.safetensors`, gate config in `config.json`)
33
+ - `lora_adapter/` — LoRA + gate weights (`modules_to_save=["text_token_gate"]`), standalone copy
34
+ - `action_head--10000_checkpoint.pt`, `proprio_projector--10000_checkpoint.pt`, `vision_backbone--10000_checkpoint.pt` (FiLM)
35
+ - `dataset_statistics.json`, `oft_training_config.json` (chunk 25, action dim 7, proprio dim 6, BOUNDS_Q99)
action_head--10000_checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:daf763567a9ae3d4e5c5f66574c6bdef6dc7be5e47b3a05656e7f322cd91b835
3
+ size 302242654
added_tokens.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ {
2
+ "<PAD>": 32000
3
+ }
config.json ADDED
@@ -0,0 +1,3192 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/workspace/kaixi/RealWorld/openvla_oft_runs/merged_public_checkpoints/oft_setting2_chunksize25_batch32_20k",
3
+ "arch_specifier": "no-align+fused-gelu-mlp",
4
+ "architectures": [
5
+ "OpenVLAForActionPrediction"
6
+ ],
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_prismatic.OpenVLAConfig",
9
+ "AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
10
+ },
11
+ "contrastive_text_tau": 0.05,
12
+ "contrastive_visual_tau": 0.1,
13
+ "gazing_mode": "self_attention",
14
+ "hf_llm_id": "meta-llama/Llama-2-7b-hf",
15
+ "image_resize_strategy": "resize-naive",
16
+ "image_sizes": [
17
+ 224,
18
+ 224
19
+ ],
20
+ "layer_gate_mode": "threshold",
21
+ "layer_gate_strength": 0.5,
22
+ "layer_gate_threshold": 0.3,
23
+ "llm_backbone_id": "llama2-7b-pure",
24
+ "llm_max_length": 2048,
25
+ "model_type": "openvla",
26
+ "n_action_bins": 256,
27
+ "norm_stats": {
28
+ "austin_buds_dataset_converted_externally_to_rlds": {
29
+ "action": {
30
+ "mask": [
31
+ true,
32
+ true,
33
+ true,
34
+ true,
35
+ true,
36
+ true,
37
+ false
38
+ ],
39
+ "max": [
40
+ 1.0,
41
+ 1.0,
42
+ 1.0,
43
+ 0.0,
44
+ 0.0,
45
+ 0.0,
46
+ 1.0
47
+ ],
48
+ "mean": [
49
+ -0.07678354531526566,
50
+ 0.0036849044263362885,
51
+ 0.05644911900162697,
52
+ 0.0,
53
+ 0.0,
54
+ 0.0,
55
+ 0.3510494828224182
56
+ ],
57
+ "min": [
58
+ -1.0,
59
+ -1.0,
60
+ -1.0,
61
+ 0.0,
62
+ 0.0,
63
+ 0.0,
64
+ 0.0
65
+ ],
66
+ "q01": [
67
+ -1.0,
68
+ -0.9599999785423279,
69
+ -0.8714285492897034,
70
+ 0.0,
71
+ 0.0,
72
+ 0.0,
73
+ 0.0
74
+ ],
75
+ "q99": [
76
+ 1.0,
77
+ 0.8600000143051147,
78
+ 1.0,
79
+ 0.0,
80
+ 0.0,
81
+ 0.0,
82
+ 1.0
83
+ ],
84
+ "std": [
85
+ 0.6367740631103516,
86
+ 0.37889179587364197,
87
+ 0.47796326875686646,
88
+ 0.0,
89
+ 0.0,
90
+ 0.0,
91
+ 0.47721168398857117
92
+ ]
93
+ },
94
+ "num_trajectories": 50,
95
+ "num_transitions": 34112,
96
+ "proprio": {
97
+ "max": [
98
+ 0.0,
99
+ 0.0,
100
+ 0.0,
101
+ 0.0,
102
+ 0.0,
103
+ 0.0,
104
+ 0.0
105
+ ],
106
+ "mean": [
107
+ 0.0,
108
+ 0.0,
109
+ 0.0,
110
+ 0.0,
111
+ 0.0,
112
+ 0.0,
113
+ 0.0
114
+ ],
115
+ "min": [
116
+ 0.0,
117
+ 0.0,
118
+ 0.0,
119
+ 0.0,
120
+ 0.0,
121
+ 0.0,
122
+ 0.0
123
+ ],
124
+ "q01": [
125
+ 0.0,
126
+ 0.0,
127
+ 0.0,
128
+ 0.0,
129
+ 0.0,
130
+ 0.0,
131
+ 0.0
132
+ ],
133
+ "q99": [
134
+ 0.0,
135
+ 0.0,
136
+ 0.0,
137
+ 0.0,
138
+ 0.0,
139
+ 0.0,
140
+ 0.0
141
+ ],
142
+ "std": [
143
+ 0.0,
144
+ 0.0,
145
+ 0.0,
146
+ 0.0,
147
+ 0.0,
148
+ 0.0,
149
+ 0.0
150
+ ]
151
+ }
152
+ },
153
+ "austin_sailor_dataset_converted_externally_to_rlds": {
154
+ "action": {
155
+ "mask": [
156
+ true,
157
+ true,
158
+ true,
159
+ true,
160
+ true,
161
+ true,
162
+ false
163
+ ],
164
+ "max": [
165
+ 1.0,
166
+ 1.0,
167
+ 1.0,
168
+ 0.0,
169
+ 0.0,
170
+ 0.375,
171
+ 1.0
172
+ ],
173
+ "mean": [
174
+ 0.011825348250567913,
175
+ 0.006461074110120535,
176
+ 0.06023626774549484,
177
+ 0.0,
178
+ 0.0,
179
+ 0.0016465914668515325,
180
+ 0.5260950326919556
181
+ ],
182
+ "min": [
183
+ -1.0,
184
+ -1.0,
185
+ -1.0,
186
+ 0.0,
187
+ 0.0,
188
+ -0.375,
189
+ 0.0
190
+ ],
191
+ "q01": [
192
+ -1.0,
193
+ -0.9828571677207947,
194
+ -0.6000000238418579,
195
+ 0.0,
196
+ 0.0,
197
+ -0.17249999940395355,
198
+ 0.0
199
+ ],
200
+ "q99": [
201
+ 1.0,
202
+ 0.9457142949104309,
203
+ 1.0,
204
+ 0.0,
205
+ 0.0,
206
+ 0.17892856895923615,
207
+ 1.0
208
+ ],
209
+ "std": [
210
+ 0.46348899602890015,
211
+ 0.41240179538726807,
212
+ 0.411862850189209,
213
+ 0.0,
214
+ 0.0,
215
+ 0.0578610822558403,
216
+ 0.49894046783447266
217
+ ]
218
+ },
219
+ "num_trajectories": 240,
220
+ "num_transitions": 353094,
221
+ "proprio": {
222
+ "max": [
223
+ 0.0,
224
+ 0.0,
225
+ 0.0,
226
+ 0.0,
227
+ 0.0,
228
+ 0.0,
229
+ 0.0
230
+ ],
231
+ "mean": [
232
+ 0.0,
233
+ 0.0,
234
+ 0.0,
235
+ 0.0,
236
+ 0.0,
237
+ 0.0,
238
+ 0.0
239
+ ],
240
+ "min": [
241
+ 0.0,
242
+ 0.0,
243
+ 0.0,
244
+ 0.0,
245
+ 0.0,
246
+ 0.0,
247
+ 0.0
248
+ ],
249
+ "q01": [
250
+ 0.0,
251
+ 0.0,
252
+ 0.0,
253
+ 0.0,
254
+ 0.0,
255
+ 0.0,
256
+ 0.0
257
+ ],
258
+ "q99": [
259
+ 0.0,
260
+ 0.0,
261
+ 0.0,
262
+ 0.0,
263
+ 0.0,
264
+ 0.0,
265
+ 0.0
266
+ ],
267
+ "std": [
268
+ 0.0,
269
+ 0.0,
270
+ 0.0,
271
+ 0.0,
272
+ 0.0,
273
+ 0.0,
274
+ 0.0
275
+ ]
276
+ }
277
+ },
278
+ "austin_sirius_dataset_converted_externally_to_rlds": {
279
+ "action": {
280
+ "mask": [
281
+ true,
282
+ true,
283
+ true,
284
+ true,
285
+ true,
286
+ true,
287
+ false
288
+ ],
289
+ "max": [
290
+ 1.0002285242080688,
291
+ 0.960608720779419,
292
+ 1.105179786682129,
293
+ 0.0,
294
+ 0.0,
295
+ 0.341785728931427,
296
+ 1.0
297
+ ],
298
+ "mean": [
299
+ 0.07747682929039001,
300
+ 0.03195561468601227,
301
+ 0.04244732856750488,
302
+ 0.0,
303
+ 0.0,
304
+ -0.01603456400334835,
305
+ 0.43260177969932556
306
+ ],
307
+ "min": [
308
+ -1.0183025598526,
309
+ -0.9800000190734863,
310
+ -0.9774575233459473,
311
+ 0.0,
312
+ 0.0,
313
+ -0.34607142210006714,
314
+ 0.0
315
+ ],
316
+ "q01": [
317
+ -0.780905865430832,
318
+ -0.5667179036140442,
319
+ -0.5254343223571777,
320
+ 0.0,
321
+ 0.0,
322
+ -0.28495091378688814,
323
+ 0.0
324
+ ],
325
+ "q99": [
326
+ 0.9569637751579284,
327
+ 0.6971374487876891,
328
+ 0.8124888157844541,
329
+ 0.0,
330
+ 0.0,
331
+ 0.1971428543329239,
332
+ 1.0
333
+ ],
334
+ "std": [
335
+ 0.3906329572200775,
336
+ 0.2998155355453491,
337
+ 0.2782271206378937,
338
+ 0.0,
339
+ 0.0,
340
+ 0.08120622485876083,
341
+ 0.49528297781944275
342
+ ]
343
+ },
344
+ "num_trajectories": 559,
345
+ "num_transitions": 279939,
346
+ "proprio": {
347
+ "max": [
348
+ 0.0,
349
+ 0.0,
350
+ 0.0,
351
+ 0.0,
352
+ 0.0,
353
+ 0.0,
354
+ 0.0
355
+ ],
356
+ "mean": [
357
+ 0.0,
358
+ 0.0,
359
+ 0.0,
360
+ 0.0,
361
+ 0.0,
362
+ 0.0,
363
+ 0.0
364
+ ],
365
+ "min": [
366
+ 0.0,
367
+ 0.0,
368
+ 0.0,
369
+ 0.0,
370
+ 0.0,
371
+ 0.0,
372
+ 0.0
373
+ ],
374
+ "q01": [
375
+ 0.0,
376
+ 0.0,
377
+ 0.0,
378
+ 0.0,
379
+ 0.0,
380
+ 0.0,
381
+ 0.0
382
+ ],
383
+ "q99": [
384
+ 0.0,
385
+ 0.0,
386
+ 0.0,
387
+ 0.0,
388
+ 0.0,
389
+ 0.0,
390
+ 0.0
391
+ ],
392
+ "std": [
393
+ 0.0,
394
+ 0.0,
395
+ 0.0,
396
+ 0.0,
397
+ 0.0,
398
+ 0.0,
399
+ 0.0
400
+ ]
401
+ }
402
+ },
403
+ "bc_z": {
404
+ "action": {
405
+ "mask": [
406
+ true,
407
+ true,
408
+ true,
409
+ true,
410
+ true,
411
+ true,
412
+ false
413
+ ],
414
+ "max": [
415
+ 0.2165454924106598,
416
+ 0.1251407265663147,
417
+ 0.10772687941789627,
418
+ 0.33544227480888367,
419
+ 0.28117990493774414,
420
+ 0.40614867210388184,
421
+ 1.0
422
+ ],
423
+ "mean": [
424
+ -0.009958467446267605,
425
+ 0.0008958321413956583,
426
+ 0.004995597992092371,
427
+ 0.00029755113064311445,
428
+ -0.008735382929444313,
429
+ -0.030693737789988518,
430
+ 0.8344562649726868
431
+ ],
432
+ "min": [
433
+ -0.1677047461271286,
434
+ -0.14630407094955444,
435
+ -0.10066790133714676,
436
+ -0.29421567916870117,
437
+ -0.32101404666900635,
438
+ -0.4635624885559082,
439
+ 0.0
440
+ ],
441
+ "q01": [
442
+ -0.09220654994249344,
443
+ -0.06456145539879798,
444
+ -0.049121275544166565,
445
+ -0.11594625547528267,
446
+ -0.14152548640966414,
447
+ -0.2251061636209488,
448
+ 0.0
449
+ ],
450
+ "q99": [
451
+ 0.07628866866230968,
452
+ 0.058019736707210584,
453
+ 0.052540797740221024,
454
+ 0.11740604028105736,
455
+ 0.11703975558280955,
456
+ 0.16729306846857078,
457
+ 1.0
458
+ ],
459
+ "std": [
460
+ 0.03053455986082554,
461
+ 0.0231423731893301,
462
+ 0.020641816779971123,
463
+ 0.04155943542718887,
464
+ 0.046427831053733826,
465
+ 0.0769818127155304,
466
+ 0.3610210120677948
467
+ ]
468
+ },
469
+ "num_trajectories": 43264,
470
+ "num_transitions": 6015535,
471
+ "proprio": {
472
+ "max": [
473
+ 0.0,
474
+ 0.0,
475
+ 0.0,
476
+ 0.0,
477
+ 0.0,
478
+ 0.0,
479
+ 0.0
480
+ ],
481
+ "mean": [
482
+ 0.0,
483
+ 0.0,
484
+ 0.0,
485
+ 0.0,
486
+ 0.0,
487
+ 0.0,
488
+ 0.0
489
+ ],
490
+ "min": [
491
+ 0.0,
492
+ 0.0,
493
+ 0.0,
494
+ 0.0,
495
+ 0.0,
496
+ 0.0,
497
+ 0.0
498
+ ],
499
+ "q01": [
500
+ 0.0,
501
+ 0.0,
502
+ 0.0,
503
+ 0.0,
504
+ 0.0,
505
+ 0.0,
506
+ 0.0
507
+ ],
508
+ "q99": [
509
+ 0.0,
510
+ 0.0,
511
+ 0.0,
512
+ 0.0,
513
+ 0.0,
514
+ 0.0,
515
+ 0.0
516
+ ],
517
+ "std": [
518
+ 0.0,
519
+ 0.0,
520
+ 0.0,
521
+ 0.0,
522
+ 0.0,
523
+ 0.0,
524
+ 0.0
525
+ ]
526
+ }
527
+ },
528
+ "berkeley_autolab_ur5": {
529
+ "action": {
530
+ "mask": [
531
+ true,
532
+ true,
533
+ true,
534
+ true,
535
+ true,
536
+ true,
537
+ false
538
+ ],
539
+ "max": [
540
+ 0.019999999552965164,
541
+ 0.019999999552965164,
542
+ 0.019999999552965164,
543
+ 0.06666667014360428,
544
+ 0.06666667014360428,
545
+ 0.06666667014360428,
546
+ 1.0
547
+ ],
548
+ "mean": [
549
+ 0.0005683620693162084,
550
+ 0.001217700308188796,
551
+ -0.0005296372692100704,
552
+ 0.00021029810886830091,
553
+ 6.0695128922816366e-05,
554
+ 0.001204986940138042,
555
+ 0.6298308372497559
556
+ ],
557
+ "min": [
558
+ -0.019999999552965164,
559
+ -0.019999999552965164,
560
+ -0.019999999552965164,
561
+ -0.06666667014360428,
562
+ -0.06666667014360428,
563
+ -0.06666667014360428,
564
+ 0.0
565
+ ],
566
+ "q01": [
567
+ -0.019999999552965164,
568
+ -0.019999999552965164,
569
+ -0.019999999552965164,
570
+ -0.02628571353852749,
571
+ -0.06666667014360428,
572
+ -0.03847619146108627,
573
+ 0.0
574
+ ],
575
+ "q99": [
576
+ 0.019999999552965164,
577
+ 0.019999999552965164,
578
+ 0.019999999552965164,
579
+ 0.031809523701667786,
580
+ 0.06666667014360428,
581
+ 0.036571428179740906,
582
+ 1.0
583
+ ],
584
+ "std": [
585
+ 0.0115329809486866,
586
+ 0.007990492507815361,
587
+ 0.009577835910022259,
588
+ 0.009432995691895485,
589
+ 0.016427582129836082,
590
+ 0.011053967289626598,
591
+ 0.48267969489097595
592
+ ]
593
+ },
594
+ "num_trajectories": 1000,
595
+ "num_transitions": 97939,
596
+ "proprio": {
597
+ "max": [
598
+ 0.0,
599
+ 0.0,
600
+ 0.0,
601
+ 0.0,
602
+ 0.0,
603
+ 0.0,
604
+ 0.0
605
+ ],
606
+ "mean": [
607
+ 0.0,
608
+ 0.0,
609
+ 0.0,
610
+ 0.0,
611
+ 0.0,
612
+ 0.0,
613
+ 0.0
614
+ ],
615
+ "min": [
616
+ 0.0,
617
+ 0.0,
618
+ 0.0,
619
+ 0.0,
620
+ 0.0,
621
+ 0.0,
622
+ 0.0
623
+ ],
624
+ "q01": [
625
+ 0.0,
626
+ 0.0,
627
+ 0.0,
628
+ 0.0,
629
+ 0.0,
630
+ 0.0,
631
+ 0.0
632
+ ],
633
+ "q99": [
634
+ 0.0,
635
+ 0.0,
636
+ 0.0,
637
+ 0.0,
638
+ 0.0,
639
+ 0.0,
640
+ 0.0
641
+ ],
642
+ "std": [
643
+ 0.0,
644
+ 0.0,
645
+ 0.0,
646
+ 0.0,
647
+ 0.0,
648
+ 0.0,
649
+ 0.0
650
+ ]
651
+ }
652
+ },
653
+ "berkeley_cable_routing": {
654
+ "action": {
655
+ "mask": [
656
+ true,
657
+ true,
658
+ true,
659
+ true,
660
+ true,
661
+ true,
662
+ false
663
+ ],
664
+ "max": [
665
+ 0.9633283019065857,
666
+ 1.0,
667
+ 1.0,
668
+ 0.0,
669
+ 0.0,
670
+ 1.0,
671
+ 0.0
672
+ ],
673
+ "mean": [
674
+ -0.07139874249696732,
675
+ 0.023609008640050888,
676
+ 0.10241943597793579,
677
+ 0.0,
678
+ 0.0,
679
+ 0.049671024084091187,
680
+ 0.0
681
+ ],
682
+ "min": [
683
+ -0.9809081554412842,
684
+ -0.9554349184036255,
685
+ -0.9994775056838989,
686
+ 0.0,
687
+ 0.0,
688
+ -1.0,
689
+ 0.0
690
+ ],
691
+ "q01": [
692
+ -0.5534318816661835,
693
+ -0.4797285574674606,
694
+ -0.5314934802055359,
695
+ 0.0,
696
+ 0.0,
697
+ -0.8855219376087189,
698
+ 0.0
699
+ ],
700
+ "q99": [
701
+ 0.42652835428714786,
702
+ 0.5000944086909298,
703
+ 0.639823433756829,
704
+ 0.0,
705
+ 0.0,
706
+ 0.984243879914284,
707
+ 0.0
708
+ ],
709
+ "std": [
710
+ 0.1815500408411026,
711
+ 0.1810990273952484,
712
+ 0.21220779418945312,
713
+ 0.0,
714
+ 0.0,
715
+ 0.3475511968135834,
716
+ 0.0
717
+ ]
718
+ },
719
+ "num_trajectories": 1647,
720
+ "num_transitions": 42328,
721
+ "proprio": {
722
+ "max": [
723
+ 0.0,
724
+ 0.0,
725
+ 0.0,
726
+ 0.0,
727
+ 0.0,
728
+ 0.0,
729
+ 0.0
730
+ ],
731
+ "mean": [
732
+ 0.0,
733
+ 0.0,
734
+ 0.0,
735
+ 0.0,
736
+ 0.0,
737
+ 0.0,
738
+ 0.0
739
+ ],
740
+ "min": [
741
+ 0.0,
742
+ 0.0,
743
+ 0.0,
744
+ 0.0,
745
+ 0.0,
746
+ 0.0,
747
+ 0.0
748
+ ],
749
+ "q01": [
750
+ 0.0,
751
+ 0.0,
752
+ 0.0,
753
+ 0.0,
754
+ 0.0,
755
+ 0.0,
756
+ 0.0
757
+ ],
758
+ "q99": [
759
+ 0.0,
760
+ 0.0,
761
+ 0.0,
762
+ 0.0,
763
+ 0.0,
764
+ 0.0,
765
+ 0.0
766
+ ],
767
+ "std": [
768
+ 0.0,
769
+ 0.0,
770
+ 0.0,
771
+ 0.0,
772
+ 0.0,
773
+ 0.0,
774
+ 0.0
775
+ ]
776
+ }
777
+ },
778
+ "berkeley_fanuc_manipulation": {
779
+ "action": {
780
+ "mask": [
781
+ true,
782
+ true,
783
+ true,
784
+ true,
785
+ true,
786
+ true,
787
+ false
788
+ ],
789
+ "max": [
790
+ 0.009999999776482582,
791
+ 0.009999999776482582,
792
+ 0.009999999776482582,
793
+ 0.03490658476948738,
794
+ 0.03490658476948738,
795
+ 0.03490658476948738,
796
+ 1.0
797
+ ],
798
+ "mean": [
799
+ 0.0007744057802483439,
800
+ -0.00031240080716088414,
801
+ -0.0015001941937953234,
802
+ -0.0007515158504247665,
803
+ -0.00015832878125365824,
804
+ 0.00014327642566058785,
805
+ 0.699295699596405
806
+ ],
807
+ "min": [
808
+ -0.009999999776482582,
809
+ -0.009999999776482582,
810
+ -0.009999999776482582,
811
+ -0.03490658476948738,
812
+ -0.03490658476948738,
813
+ -0.03490658476948738,
814
+ 0.0
815
+ ],
816
+ "q01": [
817
+ -0.009999999776482582,
818
+ -0.009999999776482582,
819
+ -0.009999999776482582,
820
+ -0.03490658476948738,
821
+ 0.0,
822
+ -0.03490658476948738,
823
+ 0.0
824
+ ],
825
+ "q99": [
826
+ 0.009999999776482582,
827
+ 0.009999999776482582,
828
+ 0.009999999776482582,
829
+ 0.03490658476948738,
830
+ 0.0,
831
+ 0.03490658476948738,
832
+ 1.0
833
+ ],
834
+ "std": [
835
+ 0.0034070091787725687,
836
+ 0.0049921851605176926,
837
+ 0.005344334989786148,
838
+ 0.00759894959628582,
839
+ 0.004081866703927517,
840
+ 0.008568956516683102,
841
+ 0.4586937427520752
842
+ ]
843
+ },
844
+ "num_trajectories": 415,
845
+ "num_transitions": 62613,
846
+ "proprio": {
847
+ "max": [
848
+ 0.0,
849
+ 0.0,
850
+ 0.0,
851
+ 0.0,
852
+ 0.0,
853
+ 0.0,
854
+ 0.0
855
+ ],
856
+ "mean": [
857
+ 0.0,
858
+ 0.0,
859
+ 0.0,
860
+ 0.0,
861
+ 0.0,
862
+ 0.0,
863
+ 0.0
864
+ ],
865
+ "min": [
866
+ 0.0,
867
+ 0.0,
868
+ 0.0,
869
+ 0.0,
870
+ 0.0,
871
+ 0.0,
872
+ 0.0
873
+ ],
874
+ "q01": [
875
+ 0.0,
876
+ 0.0,
877
+ 0.0,
878
+ 0.0,
879
+ 0.0,
880
+ 0.0,
881
+ 0.0
882
+ ],
883
+ "q99": [
884
+ 0.0,
885
+ 0.0,
886
+ 0.0,
887
+ 0.0,
888
+ 0.0,
889
+ 0.0,
890
+ 0.0
891
+ ],
892
+ "std": [
893
+ 0.0,
894
+ 0.0,
895
+ 0.0,
896
+ 0.0,
897
+ 0.0,
898
+ 0.0,
899
+ 0.0
900
+ ]
901
+ }
902
+ },
903
+ "bridge_orig": {
904
+ "action": {
905
+ "mask": [
906
+ true,
907
+ true,
908
+ true,
909
+ true,
910
+ true,
911
+ true,
912
+ false
913
+ ],
914
+ "max": [
915
+ 0.41691166162490845,
916
+ 0.25864794850349426,
917
+ 0.21218234300613403,
918
+ 3.122201919555664,
919
+ 1.8618112802505493,
920
+ 6.280478477478027,
921
+ 1.0
922
+ ],
923
+ "mean": [
924
+ 0.0002334194869035855,
925
+ 0.00013004911306779832,
926
+ -0.00012762474943883717,
927
+ -0.0001556558854645118,
928
+ -0.0004039328487124294,
929
+ 0.00023557482927571982,
930
+ 0.5764579176902771
931
+ ],
932
+ "min": [
933
+ -0.4007510244846344,
934
+ -0.13874775171279907,
935
+ -0.22553899884223938,
936
+ -3.2010786533355713,
937
+ -1.8618112802505493,
938
+ -6.279075622558594,
939
+ 0.0
940
+ ],
941
+ "q01": [
942
+ -0.02872725307941437,
943
+ -0.04170349963009357,
944
+ -0.026093858778476715,
945
+ -0.08092105075716972,
946
+ -0.09288699507713317,
947
+ -0.20718276381492615,
948
+ 0.0
949
+ ],
950
+ "q99": [
951
+ 0.028309678435325586,
952
+ 0.040855254605412394,
953
+ 0.040161586627364146,
954
+ 0.08192047759890528,
955
+ 0.07792850524187081,
956
+ 0.20382574498653397,
957
+ 1.0
958
+ ],
959
+ "std": [
960
+ 0.009765930473804474,
961
+ 0.013689135201275349,
962
+ 0.012667362578213215,
963
+ 0.028534092009067535,
964
+ 0.030637972056865692,
965
+ 0.07691419124603271,
966
+ 0.4973701536655426
967
+ ]
968
+ },
969
+ "num_trajectories": 60064,
970
+ "num_transitions": 2135463,
971
+ "proprio": {
972
+ "max": [
973
+ 0.0,
974
+ 0.0,
975
+ 0.0,
976
+ 0.0,
977
+ 0.0,
978
+ 0.0,
979
+ 0.0
980
+ ],
981
+ "mean": [
982
+ 0.0,
983
+ 0.0,
984
+ 0.0,
985
+ 0.0,
986
+ 0.0,
987
+ 0.0,
988
+ 0.0
989
+ ],
990
+ "min": [
991
+ 0.0,
992
+ 0.0,
993
+ 0.0,
994
+ 0.0,
995
+ 0.0,
996
+ 0.0,
997
+ 0.0
998
+ ],
999
+ "q01": [
1000
+ 0.0,
1001
+ 0.0,
1002
+ 0.0,
1003
+ 0.0,
1004
+ 0.0,
1005
+ 0.0,
1006
+ 0.0
1007
+ ],
1008
+ "q99": [
1009
+ 0.0,
1010
+ 0.0,
1011
+ 0.0,
1012
+ 0.0,
1013
+ 0.0,
1014
+ 0.0,
1015
+ 0.0
1016
+ ],
1017
+ "std": [
1018
+ 0.0,
1019
+ 0.0,
1020
+ 0.0,
1021
+ 0.0,
1022
+ 0.0,
1023
+ 0.0,
1024
+ 0.0
1025
+ ]
1026
+ }
1027
+ },
1028
+ "cmu_stretch": {
1029
+ "action": {
1030
+ "mask": [
1031
+ true,
1032
+ true,
1033
+ true,
1034
+ true,
1035
+ true,
1036
+ true,
1037
+ false
1038
+ ],
1039
+ "max": [
1040
+ 0.02338407188653946,
1041
+ 0.0,
1042
+ 0.023404927924275398,
1043
+ 0.0,
1044
+ 0.0,
1045
+ 0.0,
1046
+ 1.0
1047
+ ],
1048
+ "mean": [
1049
+ 0.00036304505192674696,
1050
+ 0.0,
1051
+ 0.0016466958913952112,
1052
+ 0.0,
1053
+ 0.0,
1054
+ 0.0,
1055
+ 0.3987048268318176
1056
+ ],
1057
+ "min": [
1058
+ -0.019353797659277916,
1059
+ 0.0,
1060
+ -0.02019215188920498,
1061
+ 0.0,
1062
+ 0.0,
1063
+ 0.0,
1064
+ 0.0
1065
+ ],
1066
+ "q01": [
1067
+ -0.011175686959177256,
1068
+ 0.0,
1069
+ -0.0032206363626755773,
1070
+ 0.0,
1071
+ 0.0,
1072
+ 0.0,
1073
+ 0.0
1074
+ ],
1075
+ "q99": [
1076
+ 0.014501785952597848,
1077
+ 0.0,
1078
+ 0.015056106168776728,
1079
+ 0.0,
1080
+ 0.0,
1081
+ 0.0,
1082
+ 1.0
1083
+ ],
1084
+ "std": [
1085
+ 0.004081828519701958,
1086
+ 0.0,
1087
+ 0.0037743328139185905,
1088
+ 0.0,
1089
+ 0.0,
1090
+ 0.0,
1091
+ 0.48963725566864014
1092
+ ]
1093
+ },
1094
+ "num_trajectories": 135,
1095
+ "num_transitions": 25016,
1096
+ "proprio": {
1097
+ "max": [
1098
+ 0.0,
1099
+ 0.0,
1100
+ 0.0,
1101
+ 0.0,
1102
+ 0.0,
1103
+ 0.0,
1104
+ 0.0
1105
+ ],
1106
+ "mean": [
1107
+ 0.0,
1108
+ 0.0,
1109
+ 0.0,
1110
+ 0.0,
1111
+ 0.0,
1112
+ 0.0,
1113
+ 0.0
1114
+ ],
1115
+ "min": [
1116
+ 0.0,
1117
+ 0.0,
1118
+ 0.0,
1119
+ 0.0,
1120
+ 0.0,
1121
+ 0.0,
1122
+ 0.0
1123
+ ],
1124
+ "q01": [
1125
+ 0.0,
1126
+ 0.0,
1127
+ 0.0,
1128
+ 0.0,
1129
+ 0.0,
1130
+ 0.0,
1131
+ 0.0
1132
+ ],
1133
+ "q99": [
1134
+ 0.0,
1135
+ 0.0,
1136
+ 0.0,
1137
+ 0.0,
1138
+ 0.0,
1139
+ 0.0,
1140
+ 0.0
1141
+ ],
1142
+ "std": [
1143
+ 0.0,
1144
+ 0.0,
1145
+ 0.0,
1146
+ 0.0,
1147
+ 0.0,
1148
+ 0.0,
1149
+ 0.0
1150
+ ]
1151
+ }
1152
+ },
1153
+ "dlr_edan_shared_control_converted_externally_to_rlds": {
1154
+ "action": {
1155
+ "mask": [
1156
+ true,
1157
+ true,
1158
+ true,
1159
+ true,
1160
+ true,
1161
+ true,
1162
+ false
1163
+ ],
1164
+ "max": [
1165
+ 0.18991442024707794,
1166
+ 0.0739002525806427,
1167
+ 0.18064819276332855,
1168
+ 0.0866486132144928,
1169
+ 0.13464981317520142,
1170
+ 0.16910280287265778,
1171
+ 1.0
1172
+ ],
1173
+ "mean": [
1174
+ 0.006647810339927673,
1175
+ -0.0007657372043468058,
1176
+ 0.006522852927446365,
1177
+ 0.0011679717572405934,
1178
+ -0.006395625416189432,
1179
+ -0.011902998201549053,
1180
+ 0.6985887289047241
1181
+ ],
1182
+ "min": [
1183
+ -0.10054297000169754,
1184
+ -0.08427435159683228,
1185
+ -0.13533438742160797,
1186
+ -0.17556548118591309,
1187
+ -0.18485672771930695,
1188
+ -0.2680685818195343,
1189
+ 0.0
1190
+ ],
1191
+ "q01": [
1192
+ -0.02987122368067503,
1193
+ -0.06013262912631035,
1194
+ -0.08286409199237824,
1195
+ -0.05924444157630205,
1196
+ -0.15986866518855095,
1197
+ -0.15636983573436739,
1198
+ 0.0
1199
+ ],
1200
+ "q99": [
1201
+ 0.08832092039287087,
1202
+ 0.042126184627413736,
1203
+ 0.11311905644834042,
1204
+ 0.0643695573508739,
1205
+ 0.03941855944693088,
1206
+ 0.156646853685379,
1207
+ 1.0
1208
+ ],
1209
+ "std": [
1210
+ 0.021393608301877975,
1211
+ 0.01814231649041176,
1212
+ 0.03374375030398369,
1213
+ 0.01743541844189167,
1214
+ 0.03394376486539841,
1215
+ 0.04641875624656677,
1216
+ 0.4588589072227478
1217
+ ]
1218
+ },
1219
+ "num_trajectories": 104,
1220
+ "num_transitions": 8928,
1221
+ "proprio": {
1222
+ "max": [
1223
+ 0.0,
1224
+ 0.0,
1225
+ 0.0,
1226
+ 0.0,
1227
+ 0.0,
1228
+ 0.0,
1229
+ 0.0
1230
+ ],
1231
+ "mean": [
1232
+ 0.0,
1233
+ 0.0,
1234
+ 0.0,
1235
+ 0.0,
1236
+ 0.0,
1237
+ 0.0,
1238
+ 0.0
1239
+ ],
1240
+ "min": [
1241
+ 0.0,
1242
+ 0.0,
1243
+ 0.0,
1244
+ 0.0,
1245
+ 0.0,
1246
+ 0.0,
1247
+ 0.0
1248
+ ],
1249
+ "q01": [
1250
+ 0.0,
1251
+ 0.0,
1252
+ 0.0,
1253
+ 0.0,
1254
+ 0.0,
1255
+ 0.0,
1256
+ 0.0
1257
+ ],
1258
+ "q99": [
1259
+ 0.0,
1260
+ 0.0,
1261
+ 0.0,
1262
+ 0.0,
1263
+ 0.0,
1264
+ 0.0,
1265
+ 0.0
1266
+ ],
1267
+ "std": [
1268
+ 0.0,
1269
+ 0.0,
1270
+ 0.0,
1271
+ 0.0,
1272
+ 0.0,
1273
+ 0.0,
1274
+ 0.0
1275
+ ]
1276
+ }
1277
+ },
1278
+ "dobbe": {
1279
+ "action": {
1280
+ "mask": [
1281
+ true,
1282
+ true,
1283
+ true,
1284
+ true,
1285
+ true,
1286
+ true,
1287
+ false
1288
+ ],
1289
+ "max": [
1290
+ 38.590423583984375,
1291
+ 17.932697296142578,
1292
+ 4.843764305114746,
1293
+ 1.4372116327285767,
1294
+ 0.4340403974056244,
1295
+ 1.2057193517684937,
1296
+ 0.9998947381973267
1297
+ ],
1298
+ "mean": [
1299
+ -0.0001120665911003016,
1300
+ 0.0011229600058868527,
1301
+ -0.00010194431524723768,
1302
+ -7.371398532995954e-05,
1303
+ -0.00067531579406932,
1304
+ -5.6643435527803376e-05,
1305
+ 0.6318281888961792
1306
+ ],
1307
+ "min": [
1308
+ -5.700923442840576,
1309
+ -21.605947494506836,
1310
+ -123.72489929199219,
1311
+ -1.7229845523834229,
1312
+ -0.4998578727245331,
1313
+ -0.8867913484573364,
1314
+ 1.4196479014572105e-06
1315
+ ],
1316
+ "q01": [
1317
+ -0.01119564864784479,
1318
+ -0.014266146533191203,
1319
+ -0.0071747214533388615,
1320
+ -0.009444301575422287,
1321
+ -0.03990109823644161,
1322
+ -0.017422311007976532,
1323
+ 4.003279136668425e-05
1324
+ ],
1325
+ "q99": [
1326
+ 0.01015154086053368,
1327
+ 0.017181577533483497,
1328
+ 0.007216989761218411,
1329
+ 0.010380979906767595,
1330
+ 0.03556173853576176,
1331
+ 0.018032474815845446,
1332
+ 0.9982578039169312
1333
+ ],
1334
+ "std": [
1335
+ 0.04264938458800316,
1336
+ 0.04428559169173241,
1337
+ 0.12224084138870239,
1338
+ 0.005388413090258837,
1339
+ 0.011246449314057827,
1340
+ 0.006287882570177317,
1341
+ 0.39732322096824646
1342
+ ]
1343
+ },
1344
+ "num_trajectories": 5208,
1345
+ "num_transitions": 1139911,
1346
+ "proprio": {
1347
+ "max": [
1348
+ 0.0,
1349
+ 0.0,
1350
+ 0.0,
1351
+ 0.0,
1352
+ 0.0,
1353
+ 0.0,
1354
+ 0.0
1355
+ ],
1356
+ "mean": [
1357
+ 0.0,
1358
+ 0.0,
1359
+ 0.0,
1360
+ 0.0,
1361
+ 0.0,
1362
+ 0.0,
1363
+ 0.0
1364
+ ],
1365
+ "min": [
1366
+ 0.0,
1367
+ 0.0,
1368
+ 0.0,
1369
+ 0.0,
1370
+ 0.0,
1371
+ 0.0,
1372
+ 0.0
1373
+ ],
1374
+ "q01": [
1375
+ 0.0,
1376
+ 0.0,
1377
+ 0.0,
1378
+ 0.0,
1379
+ 0.0,
1380
+ 0.0,
1381
+ 0.0
1382
+ ],
1383
+ "q99": [
1384
+ 0.0,
1385
+ 0.0,
1386
+ 0.0,
1387
+ 0.0,
1388
+ 0.0,
1389
+ 0.0,
1390
+ 0.0
1391
+ ],
1392
+ "std": [
1393
+ 0.0,
1394
+ 0.0,
1395
+ 0.0,
1396
+ 0.0,
1397
+ 0.0,
1398
+ 0.0,
1399
+ 0.0
1400
+ ]
1401
+ }
1402
+ },
1403
+ "fmb_dataset": {
1404
+ "action": {
1405
+ "mask": [
1406
+ true,
1407
+ true,
1408
+ true,
1409
+ true,
1410
+ true,
1411
+ true,
1412
+ false
1413
+ ],
1414
+ "max": [
1415
+ 1.399999976158142,
1416
+ 1.0,
1417
+ 1.399999976158142,
1418
+ 1.0,
1419
+ 1.0,
1420
+ 1.0,
1421
+ 1.0
1422
+ ],
1423
+ "mean": [
1424
+ 0.059029702097177505,
1425
+ -0.06476633995771408,
1426
+ -0.09787475317716599,
1427
+ 0.004325388930737972,
1428
+ 0.00028963794466108084,
1429
+ -0.04457257315516472,
1430
+ 0.7336440086364746
1431
+ ],
1432
+ "min": [
1433
+ -1.399999976158142,
1434
+ -1.399999976158142,
1435
+ -1.0,
1436
+ -1.0,
1437
+ -1.0,
1438
+ -1.0,
1439
+ 0.0
1440
+ ],
1441
+ "q01": [
1442
+ -0.8257142901420593,
1443
+ -1.399999976158142,
1444
+ -1.0,
1445
+ -1.0,
1446
+ -0.3028571307659149,
1447
+ -1.0,
1448
+ 0.0
1449
+ ],
1450
+ "q99": [
1451
+ 1.0,
1452
+ 0.5257142782211304,
1453
+ 1.0,
1454
+ 1.0,
1455
+ 0.3400000035762787,
1456
+ 1.0,
1457
+ 1.0
1458
+ ],
1459
+ "std": [
1460
+ 0.28809213638305664,
1461
+ 0.2820415794849396,
1462
+ 0.4626740515232086,
1463
+ 0.3266514539718628,
1464
+ 0.10842999070882797,
1465
+ 0.3440099358558655,
1466
+ 0.4435282051563263
1467
+ ]
1468
+ },
1469
+ "num_trajectories": 8612,
1470
+ "num_transitions": 1137459,
1471
+ "proprio": {
1472
+ "max": [
1473
+ 0.0,
1474
+ 0.0,
1475
+ 0.0,
1476
+ 0.0,
1477
+ 0.0,
1478
+ 0.0,
1479
+ 0.0
1480
+ ],
1481
+ "mean": [
1482
+ 0.0,
1483
+ 0.0,
1484
+ 0.0,
1485
+ 0.0,
1486
+ 0.0,
1487
+ 0.0,
1488
+ 0.0
1489
+ ],
1490
+ "min": [
1491
+ 0.0,
1492
+ 0.0,
1493
+ 0.0,
1494
+ 0.0,
1495
+ 0.0,
1496
+ 0.0,
1497
+ 0.0
1498
+ ],
1499
+ "q01": [
1500
+ 0.0,
1501
+ 0.0,
1502
+ 0.0,
1503
+ 0.0,
1504
+ 0.0,
1505
+ 0.0,
1506
+ 0.0
1507
+ ],
1508
+ "q99": [
1509
+ 0.0,
1510
+ 0.0,
1511
+ 0.0,
1512
+ 0.0,
1513
+ 0.0,
1514
+ 0.0,
1515
+ 0.0
1516
+ ],
1517
+ "std": [
1518
+ 0.0,
1519
+ 0.0,
1520
+ 0.0,
1521
+ 0.0,
1522
+ 0.0,
1523
+ 0.0,
1524
+ 0.0
1525
+ ]
1526
+ }
1527
+ },
1528
+ "fractal20220817_data": {
1529
+ "action": {
1530
+ "mask": [
1531
+ true,
1532
+ true,
1533
+ true,
1534
+ true,
1535
+ true,
1536
+ true,
1537
+ false
1538
+ ],
1539
+ "max": [
1540
+ 2.9984593391418457,
1541
+ 22.09052848815918,
1542
+ 2.7507524490356445,
1543
+ 1.570636510848999,
1544
+ 1.5321086645126343,
1545
+ 1.5691522359848022,
1546
+ 1.0
1547
+ ],
1548
+ "mean": [
1549
+ 0.006987582892179489,
1550
+ 0.006265917327255011,
1551
+ -0.01262515690177679,
1552
+ 0.04333311319351196,
1553
+ -0.005756212864071131,
1554
+ 0.0009130256366916001,
1555
+ 0.5354204773902893
1556
+ ],
1557
+ "min": [
1558
+ -2.0204520225524902,
1559
+ -5.497899532318115,
1560
+ -2.031663417816162,
1561
+ -1.569917917251587,
1562
+ -1.569892168045044,
1563
+ -1.570419430732727,
1564
+ 0.0
1565
+ ],
1566
+ "q01": [
1567
+ -0.22453527510166169,
1568
+ -0.14820013284683228,
1569
+ -0.231589707583189,
1570
+ -0.3517994859814644,
1571
+ -0.4193011274933815,
1572
+ -0.43643461108207704,
1573
+ 0.0
1574
+ ],
1575
+ "q99": [
1576
+ 0.17824687153100965,
1577
+ 0.14938379630446405,
1578
+ 0.21842354819178575,
1579
+ 0.5892666035890578,
1580
+ 0.35272657424211445,
1581
+ 0.44796681255102094,
1582
+ 1.0
1583
+ ],
1584
+ "std": [
1585
+ 0.0692116990685463,
1586
+ 0.05970962345600128,
1587
+ 0.07353084534406662,
1588
+ 0.15610496699810028,
1589
+ 0.13164450228214264,
1590
+ 0.14593800902366638,
1591
+ 0.497110515832901
1592
+ ]
1593
+ },
1594
+ "num_trajectories": 87212,
1595
+ "num_transitions": 3786400,
1596
+ "proprio": {
1597
+ "max": [
1598
+ 0.0,
1599
+ 0.0,
1600
+ 0.0,
1601
+ 0.0,
1602
+ 0.0,
1603
+ 0.0,
1604
+ 0.0
1605
+ ],
1606
+ "mean": [
1607
+ 0.0,
1608
+ 0.0,
1609
+ 0.0,
1610
+ 0.0,
1611
+ 0.0,
1612
+ 0.0,
1613
+ 0.0
1614
+ ],
1615
+ "min": [
1616
+ 0.0,
1617
+ 0.0,
1618
+ 0.0,
1619
+ 0.0,
1620
+ 0.0,
1621
+ 0.0,
1622
+ 0.0
1623
+ ],
1624
+ "q01": [
1625
+ 0.0,
1626
+ 0.0,
1627
+ 0.0,
1628
+ 0.0,
1629
+ 0.0,
1630
+ 0.0,
1631
+ 0.0
1632
+ ],
1633
+ "q99": [
1634
+ 0.0,
1635
+ 0.0,
1636
+ 0.0,
1637
+ 0.0,
1638
+ 0.0,
1639
+ 0.0,
1640
+ 0.0
1641
+ ],
1642
+ "std": [
1643
+ 0.0,
1644
+ 0.0,
1645
+ 0.0,
1646
+ 0.0,
1647
+ 0.0,
1648
+ 0.0,
1649
+ 0.0
1650
+ ]
1651
+ }
1652
+ },
1653
+ "furniture_bench_dataset_converted_externally_to_rlds": {
1654
+ "action": {
1655
+ "mask": [
1656
+ true,
1657
+ true,
1658
+ true,
1659
+ true,
1660
+ true,
1661
+ true,
1662
+ false
1663
+ ],
1664
+ "max": [
1665
+ 0.10000000149011612,
1666
+ 0.10000000149011612,
1667
+ 0.10000000149011612,
1668
+ 0.8651833534240723,
1669
+ 1.0909736156463623,
1670
+ 2.863185405731201,
1671
+ 1.0
1672
+ ],
1673
+ "mean": [
1674
+ 0.00014610752987209707,
1675
+ 0.0010830952087417245,
1676
+ 0.0006224989192560315,
1677
+ -0.003303206292912364,
1678
+ -0.0026880695950239897,
1679
+ 0.018242603167891502,
1680
+ 0.48854944109916687
1681
+ ],
1682
+ "min": [
1683
+ -0.10495579987764359,
1684
+ -0.10939455777406693,
1685
+ -0.10000000149011612,
1686
+ -0.971906840801239,
1687
+ -1.0475432872772217,
1688
+ -3.06000018119812,
1689
+ 0.0
1690
+ ],
1691
+ "q01": [
1692
+ -0.053988199681043625,
1693
+ -0.05049169331789017,
1694
+ -0.032499241530895236,
1695
+ -0.1953887003660202,
1696
+ -0.41674559473991396,
1697
+ -0.8886768388748169,
1698
+ 0.0
1699
+ ],
1700
+ "q99": [
1701
+ 0.05414841488003723,
1702
+ 0.04965164884924884,
1703
+ 0.060055799782276154,
1704
+ 0.18231668293476103,
1705
+ 0.39867786407470646,
1706
+ 0.8772023963928218,
1707
+ 1.0
1708
+ ],
1709
+ "std": [
1710
+ 0.01610708422958851,
1711
+ 0.014891477301716805,
1712
+ 0.014014219865202904,
1713
+ 0.058274295181035995,
1714
+ 0.11417088657617569,
1715
+ 0.33479776978492737,
1716
+ 0.49991825222969055
1717
+ ]
1718
+ },
1719
+ "num_trajectories": 5100,
1720
+ "num_transitions": 3948057,
1721
+ "proprio": {
1722
+ "max": [
1723
+ 0.0,
1724
+ 0.0,
1725
+ 0.0,
1726
+ 0.0,
1727
+ 0.0,
1728
+ 0.0,
1729
+ 0.0
1730
+ ],
1731
+ "mean": [
1732
+ 0.0,
1733
+ 0.0,
1734
+ 0.0,
1735
+ 0.0,
1736
+ 0.0,
1737
+ 0.0,
1738
+ 0.0
1739
+ ],
1740
+ "min": [
1741
+ 0.0,
1742
+ 0.0,
1743
+ 0.0,
1744
+ 0.0,
1745
+ 0.0,
1746
+ 0.0,
1747
+ 0.0
1748
+ ],
1749
+ "q01": [
1750
+ 0.0,
1751
+ 0.0,
1752
+ 0.0,
1753
+ 0.0,
1754
+ 0.0,
1755
+ 0.0,
1756
+ 0.0
1757
+ ],
1758
+ "q99": [
1759
+ 0.0,
1760
+ 0.0,
1761
+ 0.0,
1762
+ 0.0,
1763
+ 0.0,
1764
+ 0.0,
1765
+ 0.0
1766
+ ],
1767
+ "std": [
1768
+ 0.0,
1769
+ 0.0,
1770
+ 0.0,
1771
+ 0.0,
1772
+ 0.0,
1773
+ 0.0,
1774
+ 0.0
1775
+ ]
1776
+ }
1777
+ },
1778
+ "iamlab_cmu_pickup_insert_converted_externally_to_rlds": {
1779
+ "action": {
1780
+ "mask": [
1781
+ true,
1782
+ true,
1783
+ true,
1784
+ true,
1785
+ true,
1786
+ true,
1787
+ false
1788
+ ],
1789
+ "max": [
1790
+ 0.6634981632232666,
1791
+ 0.23428471386432648,
1792
+ 0.4308285415172577,
1793
+ 3.1415927410125732,
1794
+ 0.13647015392780304,
1795
+ 3.141592502593994,
1796
+ 1.0
1797
+ ],
1798
+ "mean": [
1799
+ 0.5274372696876526,
1800
+ 0.02858201041817665,
1801
+ 0.18712575733661652,
1802
+ 1.2339589595794678,
1803
+ 0.03226623684167862,
1804
+ -1.4199490547180176,
1805
+ 0.5550631880760193
1806
+ ],
1807
+ "min": [
1808
+ 0.3071657121181488,
1809
+ -0.29754969477653503,
1810
+ 0.06578229367733002,
1811
+ -3.1415927410125732,
1812
+ -0.04584203287959099,
1813
+ -3.141592502593994,
1814
+ 0.0
1815
+ ],
1816
+ "q01": [
1817
+ 0.3148897051811218,
1818
+ -0.20317550599575043,
1819
+ 0.06785467118024827,
1820
+ -3.140952730178833,
1821
+ -0.029743434861302376,
1822
+ -3.141091251373291,
1823
+ 0.0
1824
+ ],
1825
+ "q99": [
1826
+ 0.6472805738449097,
1827
+ 0.20846802592277527,
1828
+ 0.36855655312538155,
1829
+ 3.1409926891326903,
1830
+ 0.11424950212240226,
1831
+ 3.1410969257354737,
1832
+ 1.0
1833
+ ],
1834
+ "std": [
1835
+ 0.08108345419168472,
1836
+ 0.1116757020354271,
1837
+ 0.07747554779052734,
1838
+ 2.8737246990203857,
1839
+ 0.02774704433977604,
1840
+ 2.7678682804107666,
1841
+ 0.49695101380348206
1842
+ ]
1843
+ },
1844
+ "num_trajectories": 631,
1845
+ "num_transitions": 146241,
1846
+ "proprio": {
1847
+ "max": [
1848
+ 0.0,
1849
+ 0.0,
1850
+ 0.0,
1851
+ 0.0,
1852
+ 0.0,
1853
+ 0.0,
1854
+ 0.0
1855
+ ],
1856
+ "mean": [
1857
+ 0.0,
1858
+ 0.0,
1859
+ 0.0,
1860
+ 0.0,
1861
+ 0.0,
1862
+ 0.0,
1863
+ 0.0
1864
+ ],
1865
+ "min": [
1866
+ 0.0,
1867
+ 0.0,
1868
+ 0.0,
1869
+ 0.0,
1870
+ 0.0,
1871
+ 0.0,
1872
+ 0.0
1873
+ ],
1874
+ "q01": [
1875
+ 0.0,
1876
+ 0.0,
1877
+ 0.0,
1878
+ 0.0,
1879
+ 0.0,
1880
+ 0.0,
1881
+ 0.0
1882
+ ],
1883
+ "q99": [
1884
+ 0.0,
1885
+ 0.0,
1886
+ 0.0,
1887
+ 0.0,
1888
+ 0.0,
1889
+ 0.0,
1890
+ 0.0
1891
+ ],
1892
+ "std": [
1893
+ 0.0,
1894
+ 0.0,
1895
+ 0.0,
1896
+ 0.0,
1897
+ 0.0,
1898
+ 0.0,
1899
+ 0.0
1900
+ ]
1901
+ }
1902
+ },
1903
+ "jaco_play": {
1904
+ "action": {
1905
+ "mask": [
1906
+ true,
1907
+ true,
1908
+ true,
1909
+ true,
1910
+ true,
1911
+ true,
1912
+ false
1913
+ ],
1914
+ "max": [
1915
+ 0.20000000298023224,
1916
+ 0.20000000298023224,
1917
+ 0.20000000298023224,
1918
+ 0.0,
1919
+ 0.0,
1920
+ 0.0,
1921
+ 1.0
1922
+ ],
1923
+ "mean": [
1924
+ 0.0009658430935814977,
1925
+ -0.00580078037455678,
1926
+ -0.00395062193274498,
1927
+ 0.0,
1928
+ 0.0,
1929
+ 0.0,
1930
+ 0.34934908151626587
1931
+ ],
1932
+ "min": [
1933
+ -0.20000000298023224,
1934
+ -0.20000000298023224,
1935
+ -0.20000000298023224,
1936
+ 0.0,
1937
+ 0.0,
1938
+ 0.0,
1939
+ 0.0
1940
+ ],
1941
+ "q01": [
1942
+ -0.20000000298023224,
1943
+ -0.20000000298023224,
1944
+ -0.20000000298023224,
1945
+ 0.0,
1946
+ 0.0,
1947
+ 0.0,
1948
+ 0.0
1949
+ ],
1950
+ "q99": [
1951
+ 0.20000000298023224,
1952
+ 0.20000000298023224,
1953
+ 0.20000000298023224,
1954
+ 0.0,
1955
+ 0.0,
1956
+ 0.0,
1957
+ 1.0
1958
+ ],
1959
+ "std": [
1960
+ 0.12235074490308762,
1961
+ 0.09678777307271957,
1962
+ 0.11155334860086441,
1963
+ 0.0,
1964
+ 0.0,
1965
+ 0.0,
1966
+ 0.4768252968788147
1967
+ ]
1968
+ },
1969
+ "num_trajectories": 1085,
1970
+ "num_transitions": 77965,
1971
+ "proprio": {
1972
+ "max": [
1973
+ 0.0,
1974
+ 0.0,
1975
+ 0.0,
1976
+ 0.0,
1977
+ 0.0,
1978
+ 0.0,
1979
+ 0.0
1980
+ ],
1981
+ "mean": [
1982
+ 0.0,
1983
+ 0.0,
1984
+ 0.0,
1985
+ 0.0,
1986
+ 0.0,
1987
+ 0.0,
1988
+ 0.0
1989
+ ],
1990
+ "min": [
1991
+ 0.0,
1992
+ 0.0,
1993
+ 0.0,
1994
+ 0.0,
1995
+ 0.0,
1996
+ 0.0,
1997
+ 0.0
1998
+ ],
1999
+ "q01": [
2000
+ 0.0,
2001
+ 0.0,
2002
+ 0.0,
2003
+ 0.0,
2004
+ 0.0,
2005
+ 0.0,
2006
+ 0.0
2007
+ ],
2008
+ "q99": [
2009
+ 0.0,
2010
+ 0.0,
2011
+ 0.0,
2012
+ 0.0,
2013
+ 0.0,
2014
+ 0.0,
2015
+ 0.0
2016
+ ],
2017
+ "std": [
2018
+ 0.0,
2019
+ 0.0,
2020
+ 0.0,
2021
+ 0.0,
2022
+ 0.0,
2023
+ 0.0,
2024
+ 0.0
2025
+ ]
2026
+ }
2027
+ },
2028
+ "kuka": {
2029
+ "action": {
2030
+ "mask": [
2031
+ true,
2032
+ true,
2033
+ true,
2034
+ true,
2035
+ true,
2036
+ true,
2037
+ false
2038
+ ],
2039
+ "max": [
2040
+ 0.1697135865688324,
2041
+ 0.2777623236179352,
2042
+ 0.43710532784461975,
2043
+ 0.0,
2044
+ 0.0,
2045
+ 1.9684287309646606,
2046
+ 1.0
2047
+ ],
2048
+ "mean": [
2049
+ -0.0004668905457947403,
2050
+ 0.00040138536132872105,
2051
+ -0.001280792523175478,
2052
+ 0.0,
2053
+ 0.0,
2054
+ -0.03722453489899635,
2055
+ 0.4131543040275574
2056
+ ],
2057
+ "min": [
2058
+ -0.159867063164711,
2059
+ -0.2892282009124756,
2060
+ -0.2795473635196686,
2061
+ 0.0,
2062
+ 0.0,
2063
+ -1.9875637292861938,
2064
+ 0.0
2065
+ ],
2066
+ "q01": [
2067
+ -0.06619441494345665,
2068
+ -0.08713878810405731,
2069
+ -0.15083016991615295,
2070
+ 0.0,
2071
+ 0.0,
2072
+ -0.5415697038173676,
2073
+ 0.0
2074
+ ],
2075
+ "q99": [
2076
+ 0.06601839080452929,
2077
+ 0.08732476785779003,
2078
+ 0.18168179214000715,
2079
+ 0.0,
2080
+ 0.0,
2081
+ 0.2923380345106127,
2082
+ 1.0
2083
+ ],
2084
+ "std": [
2085
+ 0.02083250693976879,
2086
+ 0.02915887162089348,
2087
+ 0.06422865390777588,
2088
+ 0.0,
2089
+ 0.0,
2090
+ 0.14224295318126678,
2091
+ 0.49086448550224304
2092
+ ]
2093
+ },
2094
+ "num_trajectories": 209880,
2095
+ "num_transitions": 2455879,
2096
+ "proprio": {
2097
+ "max": [
2098
+ 0.0,
2099
+ 0.0,
2100
+ 0.0,
2101
+ 0.0,
2102
+ 0.0,
2103
+ 0.0,
2104
+ 0.0
2105
+ ],
2106
+ "mean": [
2107
+ 0.0,
2108
+ 0.0,
2109
+ 0.0,
2110
+ 0.0,
2111
+ 0.0,
2112
+ 0.0,
2113
+ 0.0
2114
+ ],
2115
+ "min": [
2116
+ 0.0,
2117
+ 0.0,
2118
+ 0.0,
2119
+ 0.0,
2120
+ 0.0,
2121
+ 0.0,
2122
+ 0.0
2123
+ ],
2124
+ "q01": [
2125
+ 0.0,
2126
+ 0.0,
2127
+ 0.0,
2128
+ 0.0,
2129
+ 0.0,
2130
+ 0.0,
2131
+ 0.0
2132
+ ],
2133
+ "q99": [
2134
+ 0.0,
2135
+ 0.0,
2136
+ 0.0,
2137
+ 0.0,
2138
+ 0.0,
2139
+ 0.0,
2140
+ 0.0
2141
+ ],
2142
+ "std": [
2143
+ 0.0,
2144
+ 0.0,
2145
+ 0.0,
2146
+ 0.0,
2147
+ 0.0,
2148
+ 0.0,
2149
+ 0.0
2150
+ ]
2151
+ }
2152
+ },
2153
+ "nyu_franka_play_dataset_converted_externally_to_rlds": {
2154
+ "action": {
2155
+ "mask": [
2156
+ true,
2157
+ true,
2158
+ true,
2159
+ true,
2160
+ true,
2161
+ true,
2162
+ false
2163
+ ],
2164
+ "max": [
2165
+ 0.06424188613891602,
2166
+ 0.07027634978294373,
2167
+ 0.06129661202430725,
2168
+ 6.281067848205566,
2169
+ 0.1967729926109314,
2170
+ 0.26377415657043457,
2171
+ 1.0
2172
+ ],
2173
+ "mean": [
2174
+ 0.001021989737637341,
2175
+ -0.00012002651783404872,
2176
+ 0.00032894269679673016,
2177
+ 0.0015034361276775599,
2178
+ -0.002198522910475731,
2179
+ -0.001663230243138969,
2180
+ 0.7230083346366882
2181
+ ],
2182
+ "min": [
2183
+ -0.05952230095863342,
2184
+ -0.07232445478439331,
2185
+ -0.06730806827545166,
2186
+ -6.278434753417969,
2187
+ -0.21479034423828125,
2188
+ -0.3627619743347168,
2189
+ 0.0
2190
+ ],
2191
+ "q01": [
2192
+ -0.03199600875377655,
2193
+ -0.032861671447753905,
2194
+ -0.03368805110454559,
2195
+ -0.12080862045288086,
2196
+ -0.12175218224525451,
2197
+ -0.11370223641395569,
2198
+ 0.0
2199
+ ],
2200
+ "q99": [
2201
+ 0.03101520001888276,
2202
+ 0.0373908892273903,
2203
+ 0.03646374464035038,
2204
+ 0.11764093399047852,
2205
+ 0.1258920183777809,
2206
+ 0.09366151213645942,
2207
+ 1.0
2208
+ ],
2209
+ "std": [
2210
+ 0.01327415369451046,
2211
+ 0.013215910643339157,
2212
+ 0.012822109274566174,
2213
+ 0.2732451558113098,
2214
+ 0.057022541761398315,
2215
+ 0.039172880351543427,
2216
+ 0.44752755761146545
2217
+ ]
2218
+ },
2219
+ "num_trajectories": 456,
2220
+ "num_transitions": 44875,
2221
+ "proprio": {
2222
+ "max": [
2223
+ 0.0,
2224
+ 0.0,
2225
+ 0.0,
2226
+ 0.0,
2227
+ 0.0,
2228
+ 0.0,
2229
+ 0.0
2230
+ ],
2231
+ "mean": [
2232
+ 0.0,
2233
+ 0.0,
2234
+ 0.0,
2235
+ 0.0,
2236
+ 0.0,
2237
+ 0.0,
2238
+ 0.0
2239
+ ],
2240
+ "min": [
2241
+ 0.0,
2242
+ 0.0,
2243
+ 0.0,
2244
+ 0.0,
2245
+ 0.0,
2246
+ 0.0,
2247
+ 0.0
2248
+ ],
2249
+ "q01": [
2250
+ 0.0,
2251
+ 0.0,
2252
+ 0.0,
2253
+ 0.0,
2254
+ 0.0,
2255
+ 0.0,
2256
+ 0.0
2257
+ ],
2258
+ "q99": [
2259
+ 0.0,
2260
+ 0.0,
2261
+ 0.0,
2262
+ 0.0,
2263
+ 0.0,
2264
+ 0.0,
2265
+ 0.0
2266
+ ],
2267
+ "std": [
2268
+ 0.0,
2269
+ 0.0,
2270
+ 0.0,
2271
+ 0.0,
2272
+ 0.0,
2273
+ 0.0,
2274
+ 0.0
2275
+ ]
2276
+ }
2277
+ },
2278
+ "roboturk": {
2279
+ "action": {
2280
+ "mask": [
2281
+ true,
2282
+ true,
2283
+ true,
2284
+ true,
2285
+ true,
2286
+ true,
2287
+ false
2288
+ ],
2289
+ "max": [
2290
+ 0.39124172925949097,
2291
+ 0.4601028263568878,
2292
+ 0.4870833456516266,
2293
+ 1.816888689994812,
2294
+ 1.8240282535552979,
2295
+ 1.4824820756912231,
2296
+ 1.0
2297
+ ],
2298
+ "mean": [
2299
+ 0.0014448732836171985,
2300
+ -0.0015945249469950795,
2301
+ -0.0011753785656765103,
2302
+ 0.0023012510500848293,
2303
+ -0.0009382463176734746,
2304
+ -0.00011485807772260159,
2305
+ 0.5746025443077087
2306
+ ],
2307
+ "min": [
2308
+ -0.6546999216079712,
2309
+ -0.6365841031074524,
2310
+ -0.4217723608016968,
2311
+ -1.6695482730865479,
2312
+ -1.8023357391357422,
2313
+ -1.4630827903747559,
2314
+ 0.0
2315
+ ],
2316
+ "q01": [
2317
+ -0.1342635464668274,
2318
+ -0.19996687173843383,
2319
+ -0.1482972100377083,
2320
+ -0.20720748245716095,
2321
+ -0.09676413893699647,
2322
+ -0.18075634717941286,
2323
+ 0.0
2324
+ ],
2325
+ "q99": [
2326
+ 0.14956976801157001,
2327
+ 0.1805950567126275,
2328
+ 0.18841815620660796,
2329
+ 0.21615413755178453,
2330
+ 0.09457383215427405,
2331
+ 0.18543301910162005,
2332
+ 1.0
2333
+ ],
2334
+ "std": [
2335
+ 0.04935386776924133,
2336
+ 0.0635455846786499,
2337
+ 0.061164740473032,
2338
+ 0.09553450345993042,
2339
+ 0.08420111238956451,
2340
+ 0.06517903506755829,
2341
+ 0.49452081322669983
2342
+ ]
2343
+ },
2344
+ "num_trajectories": 1995,
2345
+ "num_transitions": 187507,
2346
+ "proprio": {
2347
+ "max": [
2348
+ 0.0,
2349
+ 0.0,
2350
+ 0.0,
2351
+ 0.0,
2352
+ 0.0,
2353
+ 0.0,
2354
+ 0.0
2355
+ ],
2356
+ "mean": [
2357
+ 0.0,
2358
+ 0.0,
2359
+ 0.0,
2360
+ 0.0,
2361
+ 0.0,
2362
+ 0.0,
2363
+ 0.0
2364
+ ],
2365
+ "min": [
2366
+ 0.0,
2367
+ 0.0,
2368
+ 0.0,
2369
+ 0.0,
2370
+ 0.0,
2371
+ 0.0,
2372
+ 0.0
2373
+ ],
2374
+ "q01": [
2375
+ 0.0,
2376
+ 0.0,
2377
+ 0.0,
2378
+ 0.0,
2379
+ 0.0,
2380
+ 0.0,
2381
+ 0.0
2382
+ ],
2383
+ "q99": [
2384
+ 0.0,
2385
+ 0.0,
2386
+ 0.0,
2387
+ 0.0,
2388
+ 0.0,
2389
+ 0.0,
2390
+ 0.0
2391
+ ],
2392
+ "std": [
2393
+ 0.0,
2394
+ 0.0,
2395
+ 0.0,
2396
+ 0.0,
2397
+ 0.0,
2398
+ 0.0,
2399
+ 0.0
2400
+ ]
2401
+ }
2402
+ },
2403
+ "stanford_hydra_dataset_converted_externally_to_rlds": {
2404
+ "action": {
2405
+ "mask": [
2406
+ true,
2407
+ true,
2408
+ true,
2409
+ true,
2410
+ true,
2411
+ true,
2412
+ false
2413
+ ],
2414
+ "max": [
2415
+ 0.02499854564666748,
2416
+ 0.02499903365969658,
2417
+ 0.024999922141432762,
2418
+ 0.24974457919597626,
2419
+ 0.24997030198574066,
2420
+ 0.24999946355819702,
2421
+ 1.0
2422
+ ],
2423
+ "mean": [
2424
+ 0.0007790001109242439,
2425
+ 0.00013707754260394722,
2426
+ -0.0002548607881180942,
2427
+ 0.0012903271708637476,
2428
+ -0.004751681815832853,
2429
+ 0.002692886395379901,
2430
+ 0.48855218291282654
2431
+ ],
2432
+ "min": [
2433
+ -0.024999044835567474,
2434
+ -0.024999700486660004,
2435
+ -0.02499929815530777,
2436
+ -0.24993225932121277,
2437
+ -0.2499666064977646,
2438
+ -0.2499932497739792,
2439
+ 0.0
2440
+ ],
2441
+ "q01": [
2442
+ -0.019992006458342076,
2443
+ -0.02415412735193968,
2444
+ -0.022941758055239916,
2445
+ -0.11085530579090118,
2446
+ -0.12024572037160397,
2447
+ -0.13314770206809043,
2448
+ 0.0
2449
+ ],
2450
+ "q99": [
2451
+ 0.022886231057345868,
2452
+ 0.022358838934451335,
2453
+ 0.02410089675337076,
2454
+ 0.12370114490389822,
2455
+ 0.11323311634361738,
2456
+ 0.18474749639630164,
2457
+ 1.0
2458
+ ],
2459
+ "std": [
2460
+ 0.008022161200642586,
2461
+ 0.009131459519267082,
2462
+ 0.009574338793754578,
2463
+ 0.04122216999530792,
2464
+ 0.0384303517639637,
2465
+ 0.04606688767671585,
2466
+ 0.49976691603660583
2467
+ ]
2468
+ },
2469
+ "num_trajectories": 570,
2470
+ "num_transitions": 358234,
2471
+ "proprio": {
2472
+ "max": [
2473
+ 0.0,
2474
+ 0.0,
2475
+ 0.0,
2476
+ 0.0,
2477
+ 0.0,
2478
+ 0.0,
2479
+ 0.0
2480
+ ],
2481
+ "mean": [
2482
+ 0.0,
2483
+ 0.0,
2484
+ 0.0,
2485
+ 0.0,
2486
+ 0.0,
2487
+ 0.0,
2488
+ 0.0
2489
+ ],
2490
+ "min": [
2491
+ 0.0,
2492
+ 0.0,
2493
+ 0.0,
2494
+ 0.0,
2495
+ 0.0,
2496
+ 0.0,
2497
+ 0.0
2498
+ ],
2499
+ "q01": [
2500
+ 0.0,
2501
+ 0.0,
2502
+ 0.0,
2503
+ 0.0,
2504
+ 0.0,
2505
+ 0.0,
2506
+ 0.0
2507
+ ],
2508
+ "q99": [
2509
+ 0.0,
2510
+ 0.0,
2511
+ 0.0,
2512
+ 0.0,
2513
+ 0.0,
2514
+ 0.0,
2515
+ 0.0
2516
+ ],
2517
+ "std": [
2518
+ 0.0,
2519
+ 0.0,
2520
+ 0.0,
2521
+ 0.0,
2522
+ 0.0,
2523
+ 0.0,
2524
+ 0.0
2525
+ ]
2526
+ }
2527
+ },
2528
+ "taco_play": {
2529
+ "action": {
2530
+ "mask": [
2531
+ true,
2532
+ true,
2533
+ true,
2534
+ true,
2535
+ true,
2536
+ true,
2537
+ false
2538
+ ],
2539
+ "max": [
2540
+ 1.4915844202041626,
2541
+ 2.1842432022094727,
2542
+ 2.6836395263671875,
2543
+ 5.035226821899414,
2544
+ 2.665864944458008,
2545
+ 4.250768661499023,
2546
+ 1.0
2547
+ ],
2548
+ "mean": [
2549
+ -0.003845922416076064,
2550
+ 0.009671456180512905,
2551
+ 0.012780580669641495,
2552
+ -0.005403771996498108,
2553
+ -0.009606587700545788,
2554
+ -0.002480733208358288,
2555
+ 0.4263913035392761
2556
+ ],
2557
+ "min": [
2558
+ -4.242457866668701,
2559
+ -3.192805051803589,
2560
+ -1.3371467590332031,
2561
+ -4.202683448791504,
2562
+ -2.6722638607025146,
2563
+ -3.3467135429382324,
2564
+ 0.0
2565
+ ],
2566
+ "q01": [
2567
+ -0.7106140398979186,
2568
+ -1.056944659948349,
2569
+ -0.5878450274467468,
2570
+ -0.7682853937149048,
2571
+ -0.7180147767066956,
2572
+ -1.5527938604354858,
2573
+ 0.0
2574
+ ],
2575
+ "q99": [
2576
+ 0.6482916426658629,
2577
+ 1.0051310062408447,
2578
+ 0.9480248689651489,
2579
+ 0.6926478147506714,
2580
+ 0.6351067513227462,
2581
+ 1.628010264635086,
2582
+ 1.0
2583
+ ],
2584
+ "std": [
2585
+ 0.23254038393497467,
2586
+ 0.36298269033432007,
2587
+ 0.28692901134490967,
2588
+ 0.2617705166339874,
2589
+ 0.2438892275094986,
2590
+ 0.5216503143310547,
2591
+ 0.4946896731853485
2592
+ ]
2593
+ },
2594
+ "num_trajectories": 3603,
2595
+ "num_transitions": 237798,
2596
+ "proprio": {
2597
+ "max": [
2598
+ 0.0,
2599
+ 0.0,
2600
+ 0.0,
2601
+ 0.0,
2602
+ 0.0,
2603
+ 0.0,
2604
+ 0.0
2605
+ ],
2606
+ "mean": [
2607
+ 0.0,
2608
+ 0.0,
2609
+ 0.0,
2610
+ 0.0,
2611
+ 0.0,
2612
+ 0.0,
2613
+ 0.0
2614
+ ],
2615
+ "min": [
2616
+ 0.0,
2617
+ 0.0,
2618
+ 0.0,
2619
+ 0.0,
2620
+ 0.0,
2621
+ 0.0,
2622
+ 0.0
2623
+ ],
2624
+ "q01": [
2625
+ 0.0,
2626
+ 0.0,
2627
+ 0.0,
2628
+ 0.0,
2629
+ 0.0,
2630
+ 0.0,
2631
+ 0.0
2632
+ ],
2633
+ "q99": [
2634
+ 0.0,
2635
+ 0.0,
2636
+ 0.0,
2637
+ 0.0,
2638
+ 0.0,
2639
+ 0.0,
2640
+ 0.0
2641
+ ],
2642
+ "std": [
2643
+ 0.0,
2644
+ 0.0,
2645
+ 0.0,
2646
+ 0.0,
2647
+ 0.0,
2648
+ 0.0,
2649
+ 0.0
2650
+ ]
2651
+ }
2652
+ },
2653
+ "toto": {
2654
+ "action": {
2655
+ "mask": [
2656
+ true,
2657
+ true,
2658
+ true,
2659
+ true,
2660
+ true,
2661
+ true,
2662
+ false
2663
+ ],
2664
+ "max": [
2665
+ 0.6839867234230042,
2666
+ 0.4454185664653778,
2667
+ 0.7984078526496887,
2668
+ 2.120781660079956,
2669
+ 1.371164321899414,
2670
+ 1.4118704795837402,
2671
+ 0.0
2672
+ ],
2673
+ "mean": [
2674
+ 0.38542115688323975,
2675
+ 0.007769413758069277,
2676
+ 0.3632740378379822,
2677
+ -0.6652036905288696,
2678
+ 0.1890396922826767,
2679
+ 0.03298724442720413,
2680
+ 0.0
2681
+ ],
2682
+ "min": [
2683
+ 0.09922284632921219,
2684
+ -0.5180193781852722,
2685
+ 0.13791072368621826,
2686
+ -2.635117530822754,
2687
+ -1.0734480619430542,
2688
+ -1.9282547235488892,
2689
+ 0.0
2690
+ ],
2691
+ "q01": [
2692
+ 0.1756722891330719,
2693
+ -0.3077590811252594,
2694
+ 0.235383919775486,
2695
+ -2.0908505964279174,
2696
+ -0.6191593289375306,
2697
+ -0.7488683319091797,
2698
+ 0.0
2699
+ ],
2700
+ "q99": [
2701
+ 0.6136963081359863,
2702
+ 0.33704194784164443,
2703
+ 0.6681221985816956,
2704
+ 0.7422861719131538,
2705
+ 0.7955395007133507,
2706
+ 0.740464625358582,
2707
+ 0.0
2708
+ ],
2709
+ "std": [
2710
+ 0.12211652100086212,
2711
+ 0.19378550350666046,
2712
+ 0.10178236663341522,
2713
+ 0.5725259184837341,
2714
+ 0.29884573817253113,
2715
+ 0.3259911835193634,
2716
+ 0.0
2717
+ ]
2718
+ },
2719
+ "num_trajectories": 1003,
2720
+ "num_transitions": 325699,
2721
+ "proprio": {
2722
+ "max": [
2723
+ 0.0,
2724
+ 0.0,
2725
+ 0.0,
2726
+ 0.0,
2727
+ 0.0,
2728
+ 0.0,
2729
+ 0.0
2730
+ ],
2731
+ "mean": [
2732
+ 0.0,
2733
+ 0.0,
2734
+ 0.0,
2735
+ 0.0,
2736
+ 0.0,
2737
+ 0.0,
2738
+ 0.0
2739
+ ],
2740
+ "min": [
2741
+ 0.0,
2742
+ 0.0,
2743
+ 0.0,
2744
+ 0.0,
2745
+ 0.0,
2746
+ 0.0,
2747
+ 0.0
2748
+ ],
2749
+ "q01": [
2750
+ 0.0,
2751
+ 0.0,
2752
+ 0.0,
2753
+ 0.0,
2754
+ 0.0,
2755
+ 0.0,
2756
+ 0.0
2757
+ ],
2758
+ "q99": [
2759
+ 0.0,
2760
+ 0.0,
2761
+ 0.0,
2762
+ 0.0,
2763
+ 0.0,
2764
+ 0.0,
2765
+ 0.0
2766
+ ],
2767
+ "std": [
2768
+ 0.0,
2769
+ 0.0,
2770
+ 0.0,
2771
+ 0.0,
2772
+ 0.0,
2773
+ 0.0,
2774
+ 0.0
2775
+ ]
2776
+ }
2777
+ },
2778
+ "ucsd_kitchen_dataset_converted_externally_to_rlds": {
2779
+ "action": {
2780
+ "mask": [
2781
+ true,
2782
+ true,
2783
+ true,
2784
+ true,
2785
+ true,
2786
+ true,
2787
+ false
2788
+ ],
2789
+ "max": [
2790
+ 678.0,
2791
+ 400.0,
2792
+ 507.0,
2793
+ 180.00001525878906,
2794
+ 6.000013828277588,
2795
+ 116.99998474121094,
2796
+ 1.0
2797
+ ],
2798
+ "mean": [
2799
+ 410.37567138671875,
2800
+ 116.9518814086914,
2801
+ 192.35032653808594,
2802
+ -121.22441864013672,
2803
+ -33.84893035888672,
2804
+ 50.016136169433594,
2805
+ 0.741813600063324
2806
+ ],
2807
+ "min": [
2808
+ 172.0,
2809
+ -166.0,
2810
+ -99.99999237060547,
2811
+ -180.00001525878906,
2812
+ -89.0,
2813
+ -96.00010681152344,
2814
+ 0.0
2815
+ ],
2816
+ "q01": [
2817
+ 200.00001052856445,
2818
+ -102.31004211425781,
2819
+ -94.99993370056153,
2820
+ -180.00001525878906,
2821
+ -88.00001525878906,
2822
+ -38.999977111816406,
2823
+ 0.0
2824
+ ],
2825
+ "q99": [
2826
+ 637.0,
2827
+ 368.30999999999995,
2828
+ 493.0,
2829
+ 180.00001525878906,
2830
+ 0.999983012676239,
2831
+ 105.00001525878906,
2832
+ 1.0
2833
+ ],
2834
+ "std": [
2835
+ 122.81494903564453,
2836
+ 108.8009033203125,
2837
+ 130.303466796875,
2838
+ 116.28205108642578,
2839
+ 27.621843338012695,
2840
+ 41.02094650268555,
2841
+ 0.43763357400894165
2842
+ ]
2843
+ },
2844
+ "num_trajectories": 150,
2845
+ "num_transitions": 3970,
2846
+ "proprio": {
2847
+ "max": [
2848
+ 0.0,
2849
+ 0.0,
2850
+ 0.0,
2851
+ 0.0,
2852
+ 0.0,
2853
+ 0.0,
2854
+ 0.0
2855
+ ],
2856
+ "mean": [
2857
+ 0.0,
2858
+ 0.0,
2859
+ 0.0,
2860
+ 0.0,
2861
+ 0.0,
2862
+ 0.0,
2863
+ 0.0
2864
+ ],
2865
+ "min": [
2866
+ 0.0,
2867
+ 0.0,
2868
+ 0.0,
2869
+ 0.0,
2870
+ 0.0,
2871
+ 0.0,
2872
+ 0.0
2873
+ ],
2874
+ "q01": [
2875
+ 0.0,
2876
+ 0.0,
2877
+ 0.0,
2878
+ 0.0,
2879
+ 0.0,
2880
+ 0.0,
2881
+ 0.0
2882
+ ],
2883
+ "q99": [
2884
+ 0.0,
2885
+ 0.0,
2886
+ 0.0,
2887
+ 0.0,
2888
+ 0.0,
2889
+ 0.0,
2890
+ 0.0
2891
+ ],
2892
+ "std": [
2893
+ 0.0,
2894
+ 0.0,
2895
+ 0.0,
2896
+ 0.0,
2897
+ 0.0,
2898
+ 0.0,
2899
+ 0.0
2900
+ ]
2901
+ }
2902
+ },
2903
+ "utaustin_mutex": {
2904
+ "action": {
2905
+ "mask": [
2906
+ true,
2907
+ true,
2908
+ true,
2909
+ true,
2910
+ true,
2911
+ true,
2912
+ false
2913
+ ],
2914
+ "max": [
2915
+ 1.0,
2916
+ 1.0,
2917
+ 1.0,
2918
+ 0.375,
2919
+ 0.375,
2920
+ 0.375,
2921
+ 1.0
2922
+ ],
2923
+ "mean": [
2924
+ 0.06176406890153885,
2925
+ -0.005005486309528351,
2926
+ 0.10216785222291946,
2927
+ -0.03314131125807762,
2928
+ 0.013895004987716675,
2929
+ -0.011317633092403412,
2930
+ 0.5038976669311523
2931
+ ],
2932
+ "min": [
2933
+ -1.0,
2934
+ -1.0,
2935
+ -1.0,
2936
+ -0.375,
2937
+ -0.375,
2938
+ -0.375,
2939
+ 0.0
2940
+ ],
2941
+ "q01": [
2942
+ -0.4285714328289032,
2943
+ -0.9800000190734863,
2944
+ -0.5571428537368774,
2945
+ -0.375,
2946
+ -0.15642857551574707,
2947
+ -0.335357129573822,
2948
+ 0.0
2949
+ ],
2950
+ "q99": [
2951
+ 0.5914285778999329,
2952
+ 0.9714285731315613,
2953
+ 1.0,
2954
+ 0.3278571367263794,
2955
+ 0.207857146859169,
2956
+ 0.25607141852378845,
2957
+ 1.0
2958
+ ],
2959
+ "std": [
2960
+ 0.1875014752149582,
2961
+ 0.4468473494052887,
2962
+ 0.3792876601219177,
2963
+ 0.14097853004932404,
2964
+ 0.06453701853752136,
2965
+ 0.11765272170305252,
2966
+ 0.501045286655426
2967
+ ]
2968
+ },
2969
+ "num_trajectories": 1500,
2970
+ "num_transitions": 361883,
2971
+ "proprio": {
2972
+ "max": [
2973
+ 0.0,
2974
+ 0.0,
2975
+ 0.0,
2976
+ 0.0,
2977
+ 0.0,
2978
+ 0.0,
2979
+ 0.0
2980
+ ],
2981
+ "mean": [
2982
+ 0.0,
2983
+ 0.0,
2984
+ 0.0,
2985
+ 0.0,
2986
+ 0.0,
2987
+ 0.0,
2988
+ 0.0
2989
+ ],
2990
+ "min": [
2991
+ 0.0,
2992
+ 0.0,
2993
+ 0.0,
2994
+ 0.0,
2995
+ 0.0,
2996
+ 0.0,
2997
+ 0.0
2998
+ ],
2999
+ "q01": [
3000
+ 0.0,
3001
+ 0.0,
3002
+ 0.0,
3003
+ 0.0,
3004
+ 0.0,
3005
+ 0.0,
3006
+ 0.0
3007
+ ],
3008
+ "q99": [
3009
+ 0.0,
3010
+ 0.0,
3011
+ 0.0,
3012
+ 0.0,
3013
+ 0.0,
3014
+ 0.0,
3015
+ 0.0
3016
+ ],
3017
+ "std": [
3018
+ 0.0,
3019
+ 0.0,
3020
+ 0.0,
3021
+ 0.0,
3022
+ 0.0,
3023
+ 0.0,
3024
+ 0.0
3025
+ ]
3026
+ }
3027
+ },
3028
+ "viola": {
3029
+ "action": {
3030
+ "mask": [
3031
+ true,
3032
+ true,
3033
+ true,
3034
+ true,
3035
+ true,
3036
+ true,
3037
+ false
3038
+ ],
3039
+ "max": [
3040
+ 1.0,
3041
+ 1.0,
3042
+ 1.0,
3043
+ 0.375,
3044
+ 0.36321428418159485,
3045
+ 0.375,
3046
+ 1.0
3047
+ ],
3048
+ "mean": [
3049
+ 0.04761844128370285,
3050
+ -0.029204415157437325,
3051
+ 0.05586736649274826,
3052
+ -0.002618510741740465,
3053
+ 0.006867344491183758,
3054
+ -0.01682133786380291,
3055
+ 0.7323777675628662
3056
+ ],
3057
+ "min": [
3058
+ -1.0,
3059
+ -1.0,
3060
+ -1.0,
3061
+ -0.375,
3062
+ -0.375,
3063
+ -0.375,
3064
+ 0.0
3065
+ ],
3066
+ "q01": [
3067
+ -0.9628571271896362,
3068
+ -1.0,
3069
+ -1.0,
3070
+ -0.26249998807907104,
3071
+ -0.21321429312229156,
3072
+ -0.3385714292526245,
3073
+ 0.0
3074
+ ],
3075
+ "q99": [
3076
+ 0.9114285707473755,
3077
+ 0.868571400642395,
3078
+ 1.0,
3079
+ 0.2817857265472412,
3080
+ 0.2239285707473755,
3081
+ 0.3557142913341522,
3082
+ 1.0
3083
+ ],
3084
+ "std": [
3085
+ 0.39157867431640625,
3086
+ 0.4076525568962097,
3087
+ 0.40077948570251465,
3088
+ 0.10023996233940125,
3089
+ 0.0844319611787796,
3090
+ 0.10375042259693146,
3091
+ 0.44260647892951965
3092
+ ]
3093
+ },
3094
+ "num_trajectories": 150,
3095
+ "num_transitions": 76324,
3096
+ "proprio": {
3097
+ "max": [
3098
+ 0.0,
3099
+ 0.0,
3100
+ 0.0,
3101
+ 0.0,
3102
+ 0.0,
3103
+ 0.0,
3104
+ 0.0
3105
+ ],
3106
+ "mean": [
3107
+ 0.0,
3108
+ 0.0,
3109
+ 0.0,
3110
+ 0.0,
3111
+ 0.0,
3112
+ 0.0,
3113
+ 0.0
3114
+ ],
3115
+ "min": [
3116
+ 0.0,
3117
+ 0.0,
3118
+ 0.0,
3119
+ 0.0,
3120
+ 0.0,
3121
+ 0.0,
3122
+ 0.0
3123
+ ],
3124
+ "q01": [
3125
+ 0.0,
3126
+ 0.0,
3127
+ 0.0,
3128
+ 0.0,
3129
+ 0.0,
3130
+ 0.0,
3131
+ 0.0
3132
+ ],
3133
+ "q99": [
3134
+ 0.0,
3135
+ 0.0,
3136
+ 0.0,
3137
+ 0.0,
3138
+ 0.0,
3139
+ 0.0,
3140
+ 0.0
3141
+ ],
3142
+ "std": [
3143
+ 0.0,
3144
+ 0.0,
3145
+ 0.0,
3146
+ 0.0,
3147
+ 0.0,
3148
+ 0.0,
3149
+ 0.0
3150
+ ]
3151
+ }
3152
+ }
3153
+ },
3154
+ "output_projector_states": false,
3155
+ "pad_to_multiple_of": 64,
3156
+ "pad_token_id": 32000,
3157
+ "self_attn_dim": 512,
3158
+ "self_attn_heads": 8,
3159
+ "self_attn_layers": 1,
3160
+ "text_config": {
3161
+ "model_type": "llama",
3162
+ "pad_token_id": 32000,
3163
+ "torch_dtype": "bfloat16",
3164
+ "vocab_size": 32064
3165
+ },
3166
+ "text_token_gate_budget": 0.4,
3167
+ "text_token_gate_budget_loss_weight": 0.01,
3168
+ "text_token_gate_cross_attention_dim": 512,
3169
+ "text_token_gate_cross_attention_heads": 4,
3170
+ "text_token_gate_filter_stopwords": true,
3171
+ "text_token_gate_hidden_dim": 512,
3172
+ "text_token_gate_linear_mean_penalty_weight": 0.0,
3173
+ "text_token_gate_mlp_depth": 1,
3174
+ "text_token_gate_pool_instruction_only": true,
3175
+ "text_token_gate_text_pool_hidden_dim": 128,
3176
+ "text_token_gate_text_pooling_mode": "contrastive_alignment_score",
3177
+ "text_token_gate_use_text_summary": true,
3178
+ "text_token_gate_use_vision_tokens": true,
3179
+ "timm_model_ids": [
3180
+ "vit_large_patch14_reg4_dinov2.lvd142m",
3181
+ "vit_so400m_patch14_siglip_224"
3182
+ ],
3183
+ "timm_override_act_layers": [
3184
+ null,
3185
+ null
3186
+ ],
3187
+ "torch_dtype": "bfloat16",
3188
+ "transformers_version": "4.40.1",
3189
+ "use_fused_vision_backbone": true,
3190
+ "use_text_token_gate": true,
3191
+ "vision_backbone_id": "dinosiglip-vit-so-224px"
3192
+ }
configuration_prismatic.py ADDED
@@ -0,0 +1,319 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ configuration_prismatic.py
3
+
4
+ HuggingFace-style configuration definition for Prismatic VLMs, inheriting from `transformers.PretrainedConfig`.
5
+ Default configuration specifies `siglip-224px+7b`.
6
+ """
7
+
8
+ from typing import Any, Dict, List, Optional
9
+
10
+ from transformers import PretrainedConfig
11
+ from transformers.models.auto import CONFIG_MAPPING
12
+
13
+ # === Utilities for Mapping Prismatic names to HF names ===
14
+ # fmt: off
15
+ VALID_TEXT_TOKEN_GATE_TEXT_POOLING_MODES = {"mean", "mlp", "cross_attention", "contrastive_alignment_score"}
16
+ VALID_GAZING_MODES = {"mlp", "self_attention"}
17
+ VALID_LAYER_GATE_MODES = {"none", "threshold"}
18
+
19
+ VISION_BACKBONE_TO_RESOLUTION: Dict[str, List[int]] = {
20
+ "clip-vit-l": [224], "siglip-vit-so400m": [224], "dinov2-vit-l": [224], "in1k-vit-l": [224],
21
+
22
+ "clip-vit-l-336px": [336],
23
+ "siglip-vit-so400m-384px": [384],
24
+
25
+ "dinoclip-vit-l-336px": [336, 336],
26
+ "dinosiglip-vit-so-224px": [224, 224],
27
+ "dinosiglip-vit-so-384px": [384, 384],
28
+ }
29
+ VISION_BACKBONE_TO_TIMM_ID: Dict[str, List[str]] = {
30
+ "clip-vit-l": ["vit_large_patch14_clip_224.openai"],
31
+ "clip-vit-l-336px": ["vit_large_patch14_clip_336.openai"],
32
+
33
+ "dinov2-vit-l": ["vit_large_patch14_reg4_dinov2.lvd142m"],
34
+ "in1k-vit-l": ["vit_large_patch16_224.augreg_in21k_ft_in1k"],
35
+
36
+ "siglip-vit-so400m": ["vit_so400m_patch14_siglip_224"],
37
+ "siglip-vit-so400m-384px": ["vit_so400m_patch14_siglip_384"],
38
+
39
+ "dinoclip-vit-l-336px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_large_patch14_clip_336.openai"],
40
+ "dinosiglip-vit-so-224px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_224"],
41
+ "dinosiglip-vit-so-384px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_384"],
42
+ }
43
+ TIMM_OVERRIDE_ACT_LAYER: Dict[str, List[Optional[str]]] = {
44
+ "clip-vit-l": ["quick_gelu"], "clip-vit-l-336px": ["quick_gelu"],
45
+ "dinov2-vit-l": [None], "in1k-vit-l": [None],
46
+ "siglip-vit-so400m": [None], "siglip-vit-so400m-384px": [None],
47
+ "dinoclip-vit-l-336px": [None, "quick_gelu"],
48
+ "dinosiglip-vit-so-224px": [None, None], "dinosiglip-vit-so-384px": [None, None]
49
+ }
50
+
51
+ LLM_BACKBONE_TO_HF_PATH = {
52
+ "llama2-7b-pure": "meta-llama/Llama-2-7b-hf", "llama2-13b-pure": "meta-llama/Llama-2-13b-hf",
53
+ "llama2-7b-chat": "meta-llama/Llama-2-7b-chat-hf", "llama2-13b-chat": "meta-llama/Llama-2-13b-chat-hf",
54
+
55
+ "vicuna-v15-7b": "lmsys/vicuna-7b-v1.5", "vicuna-v15-13b": "lmsys/vicuna-13b-v1.5",
56
+
57
+ "mistral-v0.1-7b-pure": "mistralai/Mistral-7B-v0.1",
58
+ "mistral-v0.1-7b-instruct": "mistralai/Mistral-7B-Instruct-v0.1",
59
+
60
+ "phi-2-3b": "microsoft/phi-2",
61
+ }
62
+ LLM_BACKBONE_TO_HF_METACLASS = {
63
+ "llama2-7b-pure": "llama", "llama2-13b-pure": "llama", "llama2-7b-chat": "llama", "llama2-13b-chat": "llama",
64
+ "vicuna-v15-7b": "llama", "vicuna-v15-13b": "llama",
65
+
66
+ "mistral-v0.1-7b-pure": "mistral", "mistral-v0.1-7b-instruct": "mistral",
67
+
68
+ "phi-2-3b": "phi",
69
+ }
70
+
71
+ VALID_VISION_BACKBONES = set(VISION_BACKBONE_TO_RESOLUTION.keys())
72
+ VALID_LLM_BACKBONES = set(LLM_BACKBONE_TO_HF_PATH)
73
+ # fmt: on
74
+
75
+
76
+ class PrismaticConfig(PretrainedConfig):
77
+ model_type: str = "prismatic"
78
+ is_composition: bool = False
79
+
80
+ def __init__(
81
+ self,
82
+ vision_backbone_id: str = "siglip-vit-so400m",
83
+ llm_backbone_id: str = "vicuna-v15-7b",
84
+ arch_specifier: str = "no-align+gelu-mlp",
85
+ use_fused_vision_backbone: Optional[bool] = None,
86
+ image_resize_strategy: str = "letterbox",
87
+ text_config: Optional[Dict[str, Any]] = None,
88
+ llm_max_length: int = 2048,
89
+ pad_token_id: int = 32000,
90
+ pad_to_multiple_of: int = 64,
91
+ output_projector_states: bool = False,
92
+ use_text_token_gate: bool = False,
93
+ text_token_gate_use_text_summary: bool = True,
94
+ text_token_gate_use_vision_tokens: bool = True,
95
+ text_token_gate_hidden_dim: int = 512,
96
+ text_token_gate_text_pooling_mode: str = "mean",
97
+ text_token_gate_text_pool_hidden_dim: int = 128,
98
+ text_token_gate_cross_attention_dim: int = 256,
99
+ text_token_gate_cross_attention_heads: int = 1,
100
+ text_token_gate_mlp_depth: int = 1,
101
+ contrastive_visual_tau: float = 0.1,
102
+ contrastive_text_tau: float = 1.0,
103
+ gazing_mode: str = "mlp",
104
+ self_attn_dim: int = 512,
105
+ self_attn_heads: int = 8,
106
+ self_attn_layers: int = 1,
107
+ layer_gate_mode: str = "none",
108
+ layer_gate_threshold: float = 0.15,
109
+ layer_gate_strength: float = 0.5,
110
+ text_token_gate_budget: float = 0.5,
111
+ text_token_gate_budget_loss_weight: float = 0.01,
112
+ text_token_gate_linear_mean_penalty_weight: float = 0.0,
113
+ text_token_gate_pool_instruction_only: bool = False,
114
+ text_token_gate_filter_stopwords: bool = False,
115
+ **kwargs: str,
116
+ ) -> None:
117
+ # Deprecated: nonzero text_token_gate_linear_mean_penalty_weight now enables the penalty directly.
118
+ kwargs.pop("text_token_gate_use_linear_mean_penalty", None)
119
+ legacy_weighted_text_pooling = kwargs.pop("text_token_gate_use_weighted_text_pooling", None)
120
+ if legacy_weighted_text_pooling is not None and text_token_gate_text_pooling_mode == "mean":
121
+ text_token_gate_text_pooling_mode = "mlp" if legacy_weighted_text_pooling else "mean"
122
+
123
+ if vision_backbone_id not in VALID_VISION_BACKBONES:
124
+ raise ValueError(f"Vision backbone `{vision_backbone_id}` not in {VALID_VISION_BACKBONES = }")
125
+
126
+ if llm_backbone_id not in VALID_LLM_BACKBONES:
127
+ raise ValueError(f"LLM backbone `{llm_backbone_id}` not in {VALID_LLM_BACKBONES = }")
128
+
129
+ if text_token_gate_text_pooling_mode not in VALID_TEXT_TOKEN_GATE_TEXT_POOLING_MODES:
130
+ raise ValueError(
131
+ "`text_token_gate_text_pooling_mode` must be one of "
132
+ f"{sorted(VALID_TEXT_TOKEN_GATE_TEXT_POOLING_MODES)}, got {text_token_gate_text_pooling_mode!r}"
133
+ )
134
+ if text_token_gate_cross_attention_dim <= 0:
135
+ raise ValueError(
136
+ "`text_token_gate_cross_attention_dim` must be positive, "
137
+ f"got {text_token_gate_cross_attention_dim}"
138
+ )
139
+ if text_token_gate_cross_attention_heads <= 0:
140
+ raise ValueError(
141
+ "`text_token_gate_cross_attention_heads` must be positive, "
142
+ f"got {text_token_gate_cross_attention_heads}"
143
+ )
144
+ if text_token_gate_cross_attention_dim % text_token_gate_cross_attention_heads != 0:
145
+ raise ValueError(
146
+ "`text_token_gate_cross_attention_dim` must be divisible by "
147
+ "`text_token_gate_cross_attention_heads`; got "
148
+ f"{text_token_gate_cross_attention_dim} and {text_token_gate_cross_attention_heads}"
149
+ )
150
+ if text_token_gate_mlp_depth <= 0:
151
+ raise ValueError(
152
+ "`text_token_gate_mlp_depth` must be positive, "
153
+ f"got {text_token_gate_mlp_depth}"
154
+ )
155
+ if contrastive_visual_tau <= 0:
156
+ raise ValueError(f"`contrastive_visual_tau` must be positive, got {contrastive_visual_tau}")
157
+ if contrastive_text_tau <= 0:
158
+ raise ValueError(f"`contrastive_text_tau` must be positive, got {contrastive_text_tau}")
159
+ if gazing_mode not in VALID_GAZING_MODES:
160
+ raise ValueError(f"`gazing_mode` must be one of {sorted(VALID_GAZING_MODES)}, got {gazing_mode!r}")
161
+ if self_attn_dim <= 0:
162
+ raise ValueError(f"`self_attn_dim` must be positive, got {self_attn_dim}")
163
+ if self_attn_heads <= 0:
164
+ raise ValueError(f"`self_attn_heads` must be positive, got {self_attn_heads}")
165
+ if self_attn_dim % self_attn_heads != 0:
166
+ raise ValueError(
167
+ "`self_attn_dim` must be divisible by `self_attn_heads`; "
168
+ f"got {self_attn_dim} and {self_attn_heads}"
169
+ )
170
+ if self_attn_layers <= 0:
171
+ raise ValueError(f"`self_attn_layers` must be positive, got {self_attn_layers}")
172
+ if layer_gate_mode not in VALID_LAYER_GATE_MODES:
173
+ raise ValueError(
174
+ f"`layer_gate_mode` must be one of {sorted(VALID_LAYER_GATE_MODES)}, got {layer_gate_mode!r}"
175
+ )
176
+ if not 0.0 <= layer_gate_threshold <= 1.0:
177
+ raise ValueError(f"`layer_gate_threshold` must be in [0, 1], got {layer_gate_threshold}")
178
+ if not 0.0 <= layer_gate_strength <= 1.0:
179
+ raise ValueError(f"`layer_gate_strength` must be in [0, 1], got {layer_gate_strength}")
180
+
181
+ # Set Prismatic Configuration Fields
182
+ self.vision_backbone_id = vision_backbone_id
183
+ self.llm_backbone_id = llm_backbone_id
184
+ self.arch_specifier = arch_specifier
185
+ self.output_projector_states = output_projector_states
186
+ self.use_text_token_gate = use_text_token_gate
187
+ self.text_token_gate_use_text_summary = text_token_gate_use_text_summary
188
+ self.text_token_gate_use_vision_tokens = text_token_gate_use_vision_tokens
189
+ self.text_token_gate_hidden_dim = text_token_gate_hidden_dim
190
+ self.text_token_gate_text_pooling_mode = text_token_gate_text_pooling_mode
191
+ self.text_token_gate_text_pool_hidden_dim = text_token_gate_text_pool_hidden_dim
192
+ self.text_token_gate_cross_attention_dim = text_token_gate_cross_attention_dim
193
+ self.text_token_gate_cross_attention_heads = text_token_gate_cross_attention_heads
194
+ self.text_token_gate_mlp_depth = text_token_gate_mlp_depth
195
+ self.contrastive_visual_tau = contrastive_visual_tau
196
+ self.contrastive_text_tau = contrastive_text_tau
197
+ self.gazing_mode = gazing_mode
198
+ self.self_attn_dim = self_attn_dim
199
+ self.self_attn_heads = self_attn_heads
200
+ self.self_attn_layers = self_attn_layers
201
+ self.layer_gate_mode = layer_gate_mode
202
+ self.layer_gate_threshold = layer_gate_threshold
203
+ self.layer_gate_strength = layer_gate_strength
204
+ self.text_token_gate_budget = text_token_gate_budget
205
+ self.text_token_gate_budget_loss_weight = text_token_gate_budget_loss_weight
206
+ self.text_token_gate_linear_mean_penalty_weight = text_token_gate_linear_mean_penalty_weight
207
+ self.text_token_gate_pool_instruction_only = text_token_gate_pool_instruction_only
208
+ self.text_token_gate_filter_stopwords = text_token_gate_filter_stopwords
209
+
210
+ # [Contract] All vision backbone parameters are lists =>> supports fused backbones with different preprocessing
211
+ self.use_fused_vision_backbone = (
212
+ use_fused_vision_backbone
213
+ if use_fused_vision_backbone is not None
214
+ else any(self.vision_backbone_id.startswith(v) for v in ["dinoclip", "dinosiglip"])
215
+ )
216
+
217
+ self.timm_model_ids = VISION_BACKBONE_TO_TIMM_ID[self.vision_backbone_id]
218
+ self.timm_override_act_layers = TIMM_OVERRIDE_ACT_LAYER[self.vision_backbone_id]
219
+ self.image_sizes = VISION_BACKBONE_TO_RESOLUTION[self.vision_backbone_id]
220
+ self.image_resize_strategy = image_resize_strategy
221
+
222
+ self.hf_llm_id = LLM_BACKBONE_TO_HF_PATH[self.llm_backbone_id]
223
+ self.llm_max_length = llm_max_length
224
+ self.pad_token_id, self.pad_to_multiple_of = pad_token_id, pad_to_multiple_of
225
+
226
+ # [IMPORTANT] HF Utilities actually look for a `text_config` field... we need to use that specific naming!
227
+ self.text_config = (
228
+ CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]](**text_config)
229
+ if text_config is not None
230
+ else CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]]()
231
+ )
232
+
233
+ # Dispatch **kwargs to super() =>> note that `pad_token_id` collides, so we pass it in here as well...
234
+ super().__init__(pad_token_id=pad_token_id, **kwargs)
235
+
236
+
237
+ class OpenVLAConfig(PrismaticConfig):
238
+ model_type: str = "openvla"
239
+
240
+ def __init__(
241
+ self,
242
+ norm_stats: Optional[Dict[str, Dict[str, Dict[str, Dict[str, List[float]]]]]] = None,
243
+ n_action_bins: int = 256,
244
+ use_text_token_gate: bool = False,
245
+ text_token_gate_use_text_summary: bool = True,
246
+ text_token_gate_use_vision_tokens: bool = True,
247
+ text_token_gate_hidden_dim: int = 512,
248
+ text_token_gate_text_pooling_mode: str = "mean",
249
+ text_token_gate_text_pool_hidden_dim: int = 128,
250
+ text_token_gate_cross_attention_dim: int = 256,
251
+ text_token_gate_cross_attention_heads: int = 1,
252
+ text_token_gate_mlp_depth: int = 1,
253
+ contrastive_visual_tau: float = 0.1,
254
+ contrastive_text_tau: float = 1.0,
255
+ gazing_mode: str = "mlp",
256
+ self_attn_dim: int = 512,
257
+ self_attn_heads: int = 8,
258
+ self_attn_layers: int = 1,
259
+ layer_gate_mode: str = "none",
260
+ layer_gate_threshold: float = 0.15,
261
+ layer_gate_strength: float = 0.5,
262
+ text_token_gate_budget: float = 0.5,
263
+ text_token_gate_budget_loss_weight: float = 0.01,
264
+ text_token_gate_linear_mean_penalty_weight: float = 0.0,
265
+ text_token_gate_pool_instruction_only: bool = False,
266
+ text_token_gate_filter_stopwords: bool = False,
267
+ **kwargs: str,
268
+ ) -> None:
269
+ self.norm_stats, self.n_action_bins = norm_stats, n_action_bins
270
+ self.use_text_token_gate = use_text_token_gate
271
+ self.text_token_gate_use_text_summary = text_token_gate_use_text_summary
272
+ self.text_token_gate_use_vision_tokens = text_token_gate_use_vision_tokens
273
+ self.text_token_gate_hidden_dim = text_token_gate_hidden_dim
274
+ self.text_token_gate_text_pooling_mode = text_token_gate_text_pooling_mode
275
+ self.text_token_gate_text_pool_hidden_dim = text_token_gate_text_pool_hidden_dim
276
+ self.text_token_gate_cross_attention_dim = text_token_gate_cross_attention_dim
277
+ self.text_token_gate_cross_attention_heads = text_token_gate_cross_attention_heads
278
+ self.text_token_gate_mlp_depth = text_token_gate_mlp_depth
279
+ self.contrastive_visual_tau = contrastive_visual_tau
280
+ self.contrastive_text_tau = contrastive_text_tau
281
+ self.gazing_mode = gazing_mode
282
+ self.self_attn_dim = self_attn_dim
283
+ self.self_attn_heads = self_attn_heads
284
+ self.self_attn_layers = self_attn_layers
285
+ self.layer_gate_mode = layer_gate_mode
286
+ self.layer_gate_threshold = layer_gate_threshold
287
+ self.layer_gate_strength = layer_gate_strength
288
+ self.text_token_gate_budget = text_token_gate_budget
289
+ self.text_token_gate_budget_loss_weight = text_token_gate_budget_loss_weight
290
+ self.text_token_gate_linear_mean_penalty_weight = text_token_gate_linear_mean_penalty_weight
291
+ self.text_token_gate_pool_instruction_only = text_token_gate_pool_instruction_only
292
+ self.text_token_gate_filter_stopwords = text_token_gate_filter_stopwords
293
+
294
+ super().__init__(
295
+ use_text_token_gate=use_text_token_gate,
296
+ text_token_gate_use_text_summary=text_token_gate_use_text_summary,
297
+ text_token_gate_use_vision_tokens=text_token_gate_use_vision_tokens,
298
+ text_token_gate_hidden_dim=text_token_gate_hidden_dim,
299
+ text_token_gate_text_pooling_mode=text_token_gate_text_pooling_mode,
300
+ text_token_gate_text_pool_hidden_dim=text_token_gate_text_pool_hidden_dim,
301
+ text_token_gate_cross_attention_dim=text_token_gate_cross_attention_dim,
302
+ text_token_gate_cross_attention_heads=text_token_gate_cross_attention_heads,
303
+ text_token_gate_mlp_depth=text_token_gate_mlp_depth,
304
+ contrastive_visual_tau=contrastive_visual_tau,
305
+ contrastive_text_tau=contrastive_text_tau,
306
+ gazing_mode=gazing_mode,
307
+ self_attn_dim=self_attn_dim,
308
+ self_attn_heads=self_attn_heads,
309
+ self_attn_layers=self_attn_layers,
310
+ layer_gate_mode=layer_gate_mode,
311
+ layer_gate_threshold=layer_gate_threshold,
312
+ layer_gate_strength=layer_gate_strength,
313
+ text_token_gate_budget=text_token_gate_budget,
314
+ text_token_gate_budget_loss_weight=text_token_gate_budget_loss_weight,
315
+ text_token_gate_linear_mean_penalty_weight=text_token_gate_linear_mean_penalty_weight,
316
+ text_token_gate_pool_instruction_only=text_token_gate_pool_instruction_only,
317
+ text_token_gate_filter_stopwords=text_token_gate_filter_stopwords,
318
+ **kwargs,
319
+ )
dataset_statistics.json ADDED
@@ -0,0 +1,121 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "utokyo_xarm_pick_and_place_converted_externally_to_rlds": {
3
+ "action": {
4
+ "mean": [
5
+ 0.03841887041926384,
6
+ 0.04613636061549187,
7
+ -0.08122718334197998,
8
+ -0.00030900296405889094,
9
+ -0.00014601423754356802,
10
+ -0.003600441152229905,
11
+ -0.025450216606259346
12
+ ],
13
+ "std": [
14
+ 0.280966192483902,
15
+ 0.30449238419532776,
16
+ 0.34746649861335754,
17
+ 0.007257120218127966,
18
+ 0.006424501072615385,
19
+ 0.011722118593752384,
20
+ 0.9997498989105225
21
+ ],
22
+ "max": [
23
+ 1.0190002918243408,
24
+ 1.0323500633239746,
25
+ 1.0379486083984375,
26
+ 0.029906008392572403,
27
+ 0.04641101136803627,
28
+ 0.06992000341415405,
29
+ 1.0
30
+ ],
31
+ "min": [
32
+ -0.9379240870475769,
33
+ -0.9918273687362671,
34
+ -1.0000030994415283,
35
+ -0.042945992201566696,
36
+ -0.03718001767992973,
37
+ -0.06834900379180908,
38
+ -1.0
39
+ ],
40
+ "q01": [
41
+ -0.6500602000951767,
42
+ -0.825181331038475,
43
+ -0.9101772505044937,
44
+ -0.023706679213792086,
45
+ -0.017115463223308324,
46
+ -0.03339464157819748,
47
+ -1.0
48
+ ],
49
+ "q99": [
50
+ 0.8789350050687786,
51
+ 0.9065195834636688,
52
+ 0.8833758324384688,
53
+ 0.019310932401567688,
54
+ 0.01998155001550912,
55
+ 0.037749346531927574,
56
+ 1.0
57
+ ],
58
+ "mask": [
59
+ true,
60
+ true,
61
+ true,
62
+ true,
63
+ true,
64
+ true,
65
+ false
66
+ ]
67
+ },
68
+ "proprio": {
69
+ "mean": [
70
+ 0.7595251798629761,
71
+ -0.3508956730365753,
72
+ -0.9540064334869385,
73
+ -0.5224575400352478,
74
+ -0.494514524936676,
75
+ 1.6446685791015625
76
+ ],
77
+ "std": [
78
+ 0.29893943667411804,
79
+ 0.34386342763900757,
80
+ 0.38689124584198,
81
+ 1.1165426969528198,
82
+ 0.46172040700912476,
83
+ 1.6312566995620728
84
+ ],
85
+ "max": [
86
+ 1.6048250198364258,
87
+ 0.6749890446662903,
88
+ -0.21576599776744843,
89
+ 1.7784061431884766,
90
+ 1.3976739645004272,
91
+ 6.283034324645996
92
+ ],
93
+ "min": [
94
+ 0.004185997880995274,
95
+ -1.3258609771728516,
96
+ -2.389573097229004,
97
+ -4.559529781341553,
98
+ -1.7154351472854614,
99
+ -1.2300701141357422
100
+ ],
101
+ "q01": [
102
+ 0.11450984716415405,
103
+ -1.084741222858429,
104
+ -2.089673328399658,
105
+ -3.8250460863113402,
106
+ -1.315826759338379,
107
+ -0.59897809445858
108
+ ],
109
+ "q99": [
110
+ 1.50598451256752,
111
+ 0.4800056433677673,
112
+ -0.28888820469379445,
113
+ 0.4419352778792373,
114
+ 0.9082479834556553,
115
+ 5.827383484840379
116
+ ]
117
+ },
118
+ "num_transitions": 13438,
119
+ "num_trajectories": 51
120
+ }
121
+ }
generation_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 1,
4
+ "eos_token_id": 2,
5
+ "pad_token_id": 32000,
6
+ "transformers_version": "4.40.1"
7
+ }
lora_adapter/README.md ADDED
@@ -0,0 +1,202 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ base_model: /workspace/kaixi/RealWorld/openvla_oft_runs/merged_public_checkpoints/oft_setting2_chunksize25_batch32_20k
3
+ library_name: peft
4
+ ---
5
+
6
+ # Model Card for Model ID
7
+
8
+ <!-- Provide a quick summary of what the model is/does. -->
9
+
10
+
11
+
12
+ ## Model Details
13
+
14
+ ### Model Description
15
+
16
+ <!-- Provide a longer summary of what this model is. -->
17
+
18
+
19
+
20
+ - **Developed by:** [More Information Needed]
21
+ - **Funded by [optional]:** [More Information Needed]
22
+ - **Shared by [optional]:** [More Information Needed]
23
+ - **Model type:** [More Information Needed]
24
+ - **Language(s) (NLP):** [More Information Needed]
25
+ - **License:** [More Information Needed]
26
+ - **Finetuned from model [optional]:** [More Information Needed]
27
+
28
+ ### Model Sources [optional]
29
+
30
+ <!-- Provide the basic links for the model. -->
31
+
32
+ - **Repository:** [More Information Needed]
33
+ - **Paper [optional]:** [More Information Needed]
34
+ - **Demo [optional]:** [More Information Needed]
35
+
36
+ ## Uses
37
+
38
+ <!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
39
+
40
+ ### Direct Use
41
+
42
+ <!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
43
+
44
+ [More Information Needed]
45
+
46
+ ### Downstream Use [optional]
47
+
48
+ <!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
49
+
50
+ [More Information Needed]
51
+
52
+ ### Out-of-Scope Use
53
+
54
+ <!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
55
+
56
+ [More Information Needed]
57
+
58
+ ## Bias, Risks, and Limitations
59
+
60
+ <!-- This section is meant to convey both technical and sociotechnical limitations. -->
61
+
62
+ [More Information Needed]
63
+
64
+ ### Recommendations
65
+
66
+ <!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
67
+
68
+ Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
69
+
70
+ ## How to Get Started with the Model
71
+
72
+ Use the code below to get started with the model.
73
+
74
+ [More Information Needed]
75
+
76
+ ## Training Details
77
+
78
+ ### Training Data
79
+
80
+ <!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
81
+
82
+ [More Information Needed]
83
+
84
+ ### Training Procedure
85
+
86
+ <!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
87
+
88
+ #### Preprocessing [optional]
89
+
90
+ [More Information Needed]
91
+
92
+
93
+ #### Training Hyperparameters
94
+
95
+ - **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
96
+
97
+ #### Speeds, Sizes, Times [optional]
98
+
99
+ <!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
100
+
101
+ [More Information Needed]
102
+
103
+ ## Evaluation
104
+
105
+ <!-- This section describes the evaluation protocols and provides the results. -->
106
+
107
+ ### Testing Data, Factors & Metrics
108
+
109
+ #### Testing Data
110
+
111
+ <!-- This should link to a Dataset Card if possible. -->
112
+
113
+ [More Information Needed]
114
+
115
+ #### Factors
116
+
117
+ <!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
118
+
119
+ [More Information Needed]
120
+
121
+ #### Metrics
122
+
123
+ <!-- These are the evaluation metrics being used, ideally with a description of why. -->
124
+
125
+ [More Information Needed]
126
+
127
+ ### Results
128
+
129
+ [More Information Needed]
130
+
131
+ #### Summary
132
+
133
+
134
+
135
+ ## Model Examination [optional]
136
+
137
+ <!-- Relevant interpretability work for the model goes here -->
138
+
139
+ [More Information Needed]
140
+
141
+ ## Environmental Impact
142
+
143
+ <!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
144
+
145
+ Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
146
+
147
+ - **Hardware Type:** [More Information Needed]
148
+ - **Hours used:** [More Information Needed]
149
+ - **Cloud Provider:** [More Information Needed]
150
+ - **Compute Region:** [More Information Needed]
151
+ - **Carbon Emitted:** [More Information Needed]
152
+
153
+ ## Technical Specifications [optional]
154
+
155
+ ### Model Architecture and Objective
156
+
157
+ [More Information Needed]
158
+
159
+ ### Compute Infrastructure
160
+
161
+ [More Information Needed]
162
+
163
+ #### Hardware
164
+
165
+ [More Information Needed]
166
+
167
+ #### Software
168
+
169
+ [More Information Needed]
170
+
171
+ ## Citation [optional]
172
+
173
+ <!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
174
+
175
+ **BibTeX:**
176
+
177
+ [More Information Needed]
178
+
179
+ **APA:**
180
+
181
+ [More Information Needed]
182
+
183
+ ## Glossary [optional]
184
+
185
+ <!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
186
+
187
+ [More Information Needed]
188
+
189
+ ## More Information [optional]
190
+
191
+ [More Information Needed]
192
+
193
+ ## Model Card Authors [optional]
194
+
195
+ [More Information Needed]
196
+
197
+ ## Model Card Contact
198
+
199
+ [More Information Needed]
200
+ ### Framework versions
201
+
202
+ - PEFT 0.11.1
lora_adapter/adapter_config.json ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "alpha_pattern": {},
3
+ "auto_mapping": {
4
+ "base_model_class": "OpenVLAForActionPrediction",
5
+ "parent_library": "transformers_modules.oft_setting2_chunksize25_batch32_20k.modeling_prismatic"
6
+ },
7
+ "base_model_name_or_path": "/workspace/kaixi/RealWorld/openvla_oft_runs/merged_public_checkpoints/oft_setting2_chunksize25_batch32_20k",
8
+ "bias": "none",
9
+ "fan_in_fan_out": false,
10
+ "inference_mode": true,
11
+ "init_lora_weights": "gaussian",
12
+ "layer_replication": null,
13
+ "layers_pattern": null,
14
+ "layers_to_transform": null,
15
+ "loftq_config": {},
16
+ "lora_alpha": 16,
17
+ "lora_dropout": 0.0,
18
+ "megatron_config": null,
19
+ "megatron_core": "megatron.core",
20
+ "modules_to_save": [
21
+ "text_token_gate"
22
+ ],
23
+ "peft_type": "LORA",
24
+ "r": 32,
25
+ "rank_pattern": {},
26
+ "revision": null,
27
+ "target_modules": [
28
+ "fc1",
29
+ "fc2",
30
+ "gate_proj",
31
+ "up_proj",
32
+ "fc3",
33
+ "v_proj",
34
+ "q_proj",
35
+ "o_proj",
36
+ "k_proj",
37
+ "proj",
38
+ "down_proj",
39
+ "qkv"
40
+ ],
41
+ "task_type": null,
42
+ "use_dora": false,
43
+ "use_rslora": false
44
+ }
lora_adapter/adapter_model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:da912811787be529121cad1f8a523ff9c1c848c25309946a8ea65115dd055365
3
+ size 233848066
model-00001-of-00004.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:947495c22d795f71f17d5017d1209f3649d321f5b920103a51cf6a3c19de6d95
3
+ size 4925122448
model-00002-of-00004.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f7ea983d81c2f03c373e82f78f7653042daf9f85c117e2b9d69a501058e6a932
3
+ size 4947392496
model-00003-of-00004.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:54b69ccaecd035fa4dcea1034aaafccbfd15303c27f8e8e6d362950df383c85f
3
+ size 4947417456
model-00004-of-00004.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9da5fba11075bbf97fb8cebc00393a7ad6648f0f28f95adc936c41b772dad00f
3
+ size 277401170
model.safetensors.index.json ADDED
@@ -0,0 +1,1011 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 15097204610
4
+ },
5
+ "weight_map": {
6
+ "language_model.lm_head.weight": "model-00004-of-00004.safetensors",
7
+ "language_model.model.embed_tokens.weight": "model-00001-of-00004.safetensors",
8
+ "language_model.model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
9
+ "language_model.model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
10
+ "language_model.model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
11
+ "language_model.model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
12
+ "language_model.model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
13
+ "language_model.model.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
14
+ "language_model.model.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
15
+ "language_model.model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
16
+ "language_model.model.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
17
+ "language_model.model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
18
+ "language_model.model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
19
+ "language_model.model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
20
+ "language_model.model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
21
+ "language_model.model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
22
+ "language_model.model.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
23
+ "language_model.model.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
24
+ "language_model.model.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
25
+ "language_model.model.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
26
+ "language_model.model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
27
+ "language_model.model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
28
+ "language_model.model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
29
+ "language_model.model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
30
+ "language_model.model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
31
+ "language_model.model.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
32
+ "language_model.model.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
33
+ "language_model.model.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
34
+ "language_model.model.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
35
+ "language_model.model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
36
+ "language_model.model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
37
+ "language_model.model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
38
+ "language_model.model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
39
+ "language_model.model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
40
+ "language_model.model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
41
+ "language_model.model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
42
+ "language_model.model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
43
+ "language_model.model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
44
+ "language_model.model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
45
+ "language_model.model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
46
+ "language_model.model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
47
+ "language_model.model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
48
+ "language_model.model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
49
+ "language_model.model.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
50
+ "language_model.model.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
51
+ "language_model.model.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
52
+ "language_model.model.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
53
+ "language_model.model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
54
+ "language_model.model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
55
+ "language_model.model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
56
+ "language_model.model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
57
+ "language_model.model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
58
+ "language_model.model.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
59
+ "language_model.model.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
60
+ "language_model.model.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
61
+ "language_model.model.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
62
+ "language_model.model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
63
+ "language_model.model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
64
+ "language_model.model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
65
+ "language_model.model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
66
+ "language_model.model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
67
+ "language_model.model.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
68
+ "language_model.model.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
69
+ "language_model.model.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
70
+ "language_model.model.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
71
+ "language_model.model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
72
+ "language_model.model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
73
+ "language_model.model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
74
+ "language_model.model.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
75
+ "language_model.model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
76
+ "language_model.model.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
77
+ "language_model.model.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
78
+ "language_model.model.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
79
+ "language_model.model.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
80
+ "language_model.model.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
81
+ "language_model.model.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
82
+ "language_model.model.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
83
+ "language_model.model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
84
+ "language_model.model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
85
+ "language_model.model.layers.16.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
86
+ "language_model.model.layers.16.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
87
+ "language_model.model.layers.16.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
88
+ "language_model.model.layers.16.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
89
+ "language_model.model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
90
+ "language_model.model.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
91
+ "language_model.model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
92
+ "language_model.model.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
93
+ "language_model.model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
94
+ "language_model.model.layers.17.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
95
+ "language_model.model.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
96
+ "language_model.model.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
97
+ "language_model.model.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
98
+ "language_model.model.layers.18.input_layernorm.weight": "model-00002-of-00004.safetensors",
99
+ "language_model.model.layers.18.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
100
+ "language_model.model.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
101
+ "language_model.model.layers.18.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
102
+ "language_model.model.layers.18.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
103
+ "language_model.model.layers.18.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
104
+ "language_model.model.layers.18.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
105
+ "language_model.model.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
106
+ "language_model.model.layers.18.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
107
+ "language_model.model.layers.19.input_layernorm.weight": "model-00003-of-00004.safetensors",
108
+ "language_model.model.layers.19.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
109
+ "language_model.model.layers.19.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
110
+ "language_model.model.layers.19.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
111
+ "language_model.model.layers.19.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
112
+ "language_model.model.layers.19.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
113
+ "language_model.model.layers.19.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
114
+ "language_model.model.layers.19.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
115
+ "language_model.model.layers.19.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
116
+ "language_model.model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
117
+ "language_model.model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
118
+ "language_model.model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
119
+ "language_model.model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
120
+ "language_model.model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
121
+ "language_model.model.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
122
+ "language_model.model.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
123
+ "language_model.model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
124
+ "language_model.model.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
125
+ "language_model.model.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
126
+ "language_model.model.layers.20.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
127
+ "language_model.model.layers.20.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
128
+ "language_model.model.layers.20.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
129
+ "language_model.model.layers.20.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
130
+ "language_model.model.layers.20.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
131
+ "language_model.model.layers.20.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
132
+ "language_model.model.layers.20.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
133
+ "language_model.model.layers.20.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
134
+ "language_model.model.layers.21.input_layernorm.weight": "model-00003-of-00004.safetensors",
135
+ "language_model.model.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
136
+ "language_model.model.layers.21.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
137
+ "language_model.model.layers.21.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
138
+ "language_model.model.layers.21.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
139
+ "language_model.model.layers.21.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
140
+ "language_model.model.layers.21.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
141
+ "language_model.model.layers.21.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
142
+ "language_model.model.layers.21.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
143
+ "language_model.model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
144
+ "language_model.model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
145
+ "language_model.model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
146
+ "language_model.model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
147
+ "language_model.model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
148
+ "language_model.model.layers.22.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
149
+ "language_model.model.layers.22.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
150
+ "language_model.model.layers.22.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
151
+ "language_model.model.layers.22.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
152
+ "language_model.model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
153
+ "language_model.model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
154
+ "language_model.model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
155
+ "language_model.model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
156
+ "language_model.model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
157
+ "language_model.model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
158
+ "language_model.model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
159
+ "language_model.model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
160
+ "language_model.model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
161
+ "language_model.model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
162
+ "language_model.model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
163
+ "language_model.model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
164
+ "language_model.model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
165
+ "language_model.model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
166
+ "language_model.model.layers.24.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
167
+ "language_model.model.layers.24.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
168
+ "language_model.model.layers.24.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
169
+ "language_model.model.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
170
+ "language_model.model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
171
+ "language_model.model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
172
+ "language_model.model.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
173
+ "language_model.model.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
174
+ "language_model.model.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
175
+ "language_model.model.layers.25.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
176
+ "language_model.model.layers.25.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
177
+ "language_model.model.layers.25.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
178
+ "language_model.model.layers.25.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
179
+ "language_model.model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
180
+ "language_model.model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
181
+ "language_model.model.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
182
+ "language_model.model.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
183
+ "language_model.model.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
184
+ "language_model.model.layers.26.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
185
+ "language_model.model.layers.26.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
186
+ "language_model.model.layers.26.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
187
+ "language_model.model.layers.26.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
188
+ "language_model.model.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
189
+ "language_model.model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
190
+ "language_model.model.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
191
+ "language_model.model.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
192
+ "language_model.model.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
193
+ "language_model.model.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
194
+ "language_model.model.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
195
+ "language_model.model.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
196
+ "language_model.model.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
197
+ "language_model.model.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
198
+ "language_model.model.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
199
+ "language_model.model.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
200
+ "language_model.model.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
201
+ "language_model.model.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
202
+ "language_model.model.layers.28.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
203
+ "language_model.model.layers.28.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
204
+ "language_model.model.layers.28.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
205
+ "language_model.model.layers.28.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
206
+ "language_model.model.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
207
+ "language_model.model.layers.29.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
208
+ "language_model.model.layers.29.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
209
+ "language_model.model.layers.29.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
210
+ "language_model.model.layers.29.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
211
+ "language_model.model.layers.29.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
212
+ "language_model.model.layers.29.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
213
+ "language_model.model.layers.29.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
214
+ "language_model.model.layers.29.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
215
+ "language_model.model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
216
+ "language_model.model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
217
+ "language_model.model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
218
+ "language_model.model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
219
+ "language_model.model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
220
+ "language_model.model.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
221
+ "language_model.model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
222
+ "language_model.model.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
223
+ "language_model.model.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
224
+ "language_model.model.layers.30.input_layernorm.weight": "model-00003-of-00004.safetensors",
225
+ "language_model.model.layers.30.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
226
+ "language_model.model.layers.30.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
227
+ "language_model.model.layers.30.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
228
+ "language_model.model.layers.30.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
229
+ "language_model.model.layers.30.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
230
+ "language_model.model.layers.30.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
231
+ "language_model.model.layers.30.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
232
+ "language_model.model.layers.30.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
233
+ "language_model.model.layers.31.input_layernorm.weight": "model-00003-of-00004.safetensors",
234
+ "language_model.model.layers.31.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
235
+ "language_model.model.layers.31.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
236
+ "language_model.model.layers.31.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
237
+ "language_model.model.layers.31.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
238
+ "language_model.model.layers.31.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
239
+ "language_model.model.layers.31.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
240
+ "language_model.model.layers.31.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
241
+ "language_model.model.layers.31.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
242
+ "language_model.model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
243
+ "language_model.model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
244
+ "language_model.model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
245
+ "language_model.model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
246
+ "language_model.model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
247
+ "language_model.model.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
248
+ "language_model.model.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
249
+ "language_model.model.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
250
+ "language_model.model.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
251
+ "language_model.model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
252
+ "language_model.model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
253
+ "language_model.model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
254
+ "language_model.model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
255
+ "language_model.model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
256
+ "language_model.model.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
257
+ "language_model.model.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
258
+ "language_model.model.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
259
+ "language_model.model.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
260
+ "language_model.model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
261
+ "language_model.model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
262
+ "language_model.model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
263
+ "language_model.model.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
264
+ "language_model.model.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
265
+ "language_model.model.layers.6.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
266
+ "language_model.model.layers.6.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
267
+ "language_model.model.layers.6.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
268
+ "language_model.model.layers.6.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
269
+ "language_model.model.layers.7.input_layernorm.weight": "model-00002-of-00004.safetensors",
270
+ "language_model.model.layers.7.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
271
+ "language_model.model.layers.7.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
272
+ "language_model.model.layers.7.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
273
+ "language_model.model.layers.7.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
274
+ "language_model.model.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
275
+ "language_model.model.layers.7.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
276
+ "language_model.model.layers.7.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
277
+ "language_model.model.layers.7.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
278
+ "language_model.model.layers.8.input_layernorm.weight": "model-00002-of-00004.safetensors",
279
+ "language_model.model.layers.8.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
280
+ "language_model.model.layers.8.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
281
+ "language_model.model.layers.8.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
282
+ "language_model.model.layers.8.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
283
+ "language_model.model.layers.8.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
284
+ "language_model.model.layers.8.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
285
+ "language_model.model.layers.8.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
286
+ "language_model.model.layers.8.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
287
+ "language_model.model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
288
+ "language_model.model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
289
+ "language_model.model.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
290
+ "language_model.model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
291
+ "language_model.model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
292
+ "language_model.model.layers.9.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
293
+ "language_model.model.layers.9.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
294
+ "language_model.model.layers.9.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
295
+ "language_model.model.layers.9.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
296
+ "language_model.model.norm.weight": "model-00003-of-00004.safetensors",
297
+ "projector.fc1.bias": "model-00001-of-00004.safetensors",
298
+ "projector.fc1.weight": "model-00001-of-00004.safetensors",
299
+ "projector.fc2.bias": "model-00001-of-00004.safetensors",
300
+ "projector.fc2.weight": "model-00001-of-00004.safetensors",
301
+ "projector.fc3.bias": "model-00001-of-00004.safetensors",
302
+ "projector.fc3.weight": "model-00001-of-00004.safetensors",
303
+ "text_token_gate.context_proj.bias": "model-00004-of-00004.safetensors",
304
+ "text_token_gate.context_proj.weight": "model-00004-of-00004.safetensors",
305
+ "text_token_gate.gate_out.bias": "model-00004-of-00004.safetensors",
306
+ "text_token_gate.gate_out.weight": "model-00004-of-00004.safetensors",
307
+ "text_token_gate.gate_text_ln.bias": "model-00004-of-00004.safetensors",
308
+ "text_token_gate.gate_text_ln.weight": "model-00004-of-00004.safetensors",
309
+ "text_token_gate.gate_visual_ln.bias": "model-00004-of-00004.safetensors",
310
+ "text_token_gate.gate_visual_ln.weight": "model-00004-of-00004.safetensors",
311
+ "text_token_gate.gazing_blocks.0.attn.in_proj_bias": "model-00004-of-00004.safetensors",
312
+ "text_token_gate.gazing_blocks.0.attn.in_proj_weight": "model-00004-of-00004.safetensors",
313
+ "text_token_gate.gazing_blocks.0.attn.out_proj.bias": "model-00004-of-00004.safetensors",
314
+ "text_token_gate.gazing_blocks.0.attn.out_proj.weight": "model-00004-of-00004.safetensors",
315
+ "text_token_gate.gazing_blocks.0.attn_norm.bias": "model-00004-of-00004.safetensors",
316
+ "text_token_gate.gazing_blocks.0.attn_norm.weight": "model-00004-of-00004.safetensors",
317
+ "text_token_gate.gazing_blocks.0.ffn.0.bias": "model-00004-of-00004.safetensors",
318
+ "text_token_gate.gazing_blocks.0.ffn.0.weight": "model-00004-of-00004.safetensors",
319
+ "text_token_gate.gazing_blocks.0.ffn.2.bias": "model-00004-of-00004.safetensors",
320
+ "text_token_gate.gazing_blocks.0.ffn.2.weight": "model-00004-of-00004.safetensors",
321
+ "text_token_gate.gazing_blocks.0.ffn_norm.bias": "model-00004-of-00004.safetensors",
322
+ "text_token_gate.gazing_blocks.0.ffn_norm.weight": "model-00004-of-00004.safetensors",
323
+ "text_token_gate.gazing_output_ln.bias": "model-00004-of-00004.safetensors",
324
+ "text_token_gate.gazing_output_ln.weight": "model-00004-of-00004.safetensors",
325
+ "vision_backbone.featurizer.blocks.0.attn.proj.bias": "model-00001-of-00004.safetensors",
326
+ "vision_backbone.featurizer.blocks.0.attn.proj.weight": "model-00001-of-00004.safetensors",
327
+ "vision_backbone.featurizer.blocks.0.attn.qkv.bias": "model-00001-of-00004.safetensors",
328
+ "vision_backbone.featurizer.blocks.0.attn.qkv.weight": "model-00001-of-00004.safetensors",
329
+ "vision_backbone.featurizer.blocks.0.ls1.scale_factor": "model-00001-of-00004.safetensors",
330
+ "vision_backbone.featurizer.blocks.0.ls2.scale_factor": "model-00001-of-00004.safetensors",
331
+ "vision_backbone.featurizer.blocks.0.mlp.fc1.bias": "model-00001-of-00004.safetensors",
332
+ "vision_backbone.featurizer.blocks.0.mlp.fc1.weight": "model-00001-of-00004.safetensors",
333
+ "vision_backbone.featurizer.blocks.0.mlp.fc2.bias": "model-00001-of-00004.safetensors",
334
+ "vision_backbone.featurizer.blocks.0.mlp.fc2.weight": "model-00001-of-00004.safetensors",
335
+ "vision_backbone.featurizer.blocks.0.norm1.bias": "model-00001-of-00004.safetensors",
336
+ "vision_backbone.featurizer.blocks.0.norm1.weight": "model-00001-of-00004.safetensors",
337
+ "vision_backbone.featurizer.blocks.0.norm2.bias": "model-00001-of-00004.safetensors",
338
+ "vision_backbone.featurizer.blocks.0.norm2.weight": "model-00001-of-00004.safetensors",
339
+ "vision_backbone.featurizer.blocks.1.attn.proj.bias": "model-00001-of-00004.safetensors",
340
+ "vision_backbone.featurizer.blocks.1.attn.proj.weight": "model-00001-of-00004.safetensors",
341
+ "vision_backbone.featurizer.blocks.1.attn.qkv.bias": "model-00001-of-00004.safetensors",
342
+ "vision_backbone.featurizer.blocks.1.attn.qkv.weight": "model-00001-of-00004.safetensors",
343
+ "vision_backbone.featurizer.blocks.1.ls1.scale_factor": "model-00001-of-00004.safetensors",
344
+ "vision_backbone.featurizer.blocks.1.ls2.scale_factor": "model-00001-of-00004.safetensors",
345
+ "vision_backbone.featurizer.blocks.1.mlp.fc1.bias": "model-00001-of-00004.safetensors",
346
+ "vision_backbone.featurizer.blocks.1.mlp.fc1.weight": "model-00001-of-00004.safetensors",
347
+ "vision_backbone.featurizer.blocks.1.mlp.fc2.bias": "model-00001-of-00004.safetensors",
348
+ "vision_backbone.featurizer.blocks.1.mlp.fc2.weight": "model-00001-of-00004.safetensors",
349
+ "vision_backbone.featurizer.blocks.1.norm1.bias": "model-00001-of-00004.safetensors",
350
+ "vision_backbone.featurizer.blocks.1.norm1.weight": "model-00001-of-00004.safetensors",
351
+ "vision_backbone.featurizer.blocks.1.norm2.bias": "model-00001-of-00004.safetensors",
352
+ "vision_backbone.featurizer.blocks.1.norm2.weight": "model-00001-of-00004.safetensors",
353
+ "vision_backbone.featurizer.blocks.10.attn.proj.bias": "model-00001-of-00004.safetensors",
354
+ "vision_backbone.featurizer.blocks.10.attn.proj.weight": "model-00001-of-00004.safetensors",
355
+ "vision_backbone.featurizer.blocks.10.attn.qkv.bias": "model-00001-of-00004.safetensors",
356
+ "vision_backbone.featurizer.blocks.10.attn.qkv.weight": "model-00001-of-00004.safetensors",
357
+ "vision_backbone.featurizer.blocks.10.ls1.scale_factor": "model-00001-of-00004.safetensors",
358
+ "vision_backbone.featurizer.blocks.10.ls2.scale_factor": "model-00001-of-00004.safetensors",
359
+ "vision_backbone.featurizer.blocks.10.mlp.fc1.bias": "model-00001-of-00004.safetensors",
360
+ "vision_backbone.featurizer.blocks.10.mlp.fc1.weight": "model-00001-of-00004.safetensors",
361
+ "vision_backbone.featurizer.blocks.10.mlp.fc2.bias": "model-00001-of-00004.safetensors",
362
+ "vision_backbone.featurizer.blocks.10.mlp.fc2.weight": "model-00001-of-00004.safetensors",
363
+ "vision_backbone.featurizer.blocks.10.norm1.bias": "model-00001-of-00004.safetensors",
364
+ "vision_backbone.featurizer.blocks.10.norm1.weight": "model-00001-of-00004.safetensors",
365
+ "vision_backbone.featurizer.blocks.10.norm2.bias": "model-00001-of-00004.safetensors",
366
+ "vision_backbone.featurizer.blocks.10.norm2.weight": "model-00001-of-00004.safetensors",
367
+ "vision_backbone.featurizer.blocks.11.attn.proj.bias": "model-00001-of-00004.safetensors",
368
+ "vision_backbone.featurizer.blocks.11.attn.proj.weight": "model-00001-of-00004.safetensors",
369
+ "vision_backbone.featurizer.blocks.11.attn.qkv.bias": "model-00001-of-00004.safetensors",
370
+ "vision_backbone.featurizer.blocks.11.attn.qkv.weight": "model-00001-of-00004.safetensors",
371
+ "vision_backbone.featurizer.blocks.11.ls1.scale_factor": "model-00001-of-00004.safetensors",
372
+ "vision_backbone.featurizer.blocks.11.ls2.scale_factor": "model-00001-of-00004.safetensors",
373
+ "vision_backbone.featurizer.blocks.11.mlp.fc1.bias": "model-00001-of-00004.safetensors",
374
+ "vision_backbone.featurizer.blocks.11.mlp.fc1.weight": "model-00001-of-00004.safetensors",
375
+ "vision_backbone.featurizer.blocks.11.mlp.fc2.bias": "model-00001-of-00004.safetensors",
376
+ "vision_backbone.featurizer.blocks.11.mlp.fc2.weight": "model-00001-of-00004.safetensors",
377
+ "vision_backbone.featurizer.blocks.11.norm1.bias": "model-00001-of-00004.safetensors",
378
+ "vision_backbone.featurizer.blocks.11.norm1.weight": "model-00001-of-00004.safetensors",
379
+ "vision_backbone.featurizer.blocks.11.norm2.bias": "model-00001-of-00004.safetensors",
380
+ "vision_backbone.featurizer.blocks.11.norm2.weight": "model-00001-of-00004.safetensors",
381
+ "vision_backbone.featurizer.blocks.12.attn.proj.bias": "model-00001-of-00004.safetensors",
382
+ "vision_backbone.featurizer.blocks.12.attn.proj.weight": "model-00001-of-00004.safetensors",
383
+ "vision_backbone.featurizer.blocks.12.attn.qkv.bias": "model-00001-of-00004.safetensors",
384
+ "vision_backbone.featurizer.blocks.12.attn.qkv.weight": "model-00001-of-00004.safetensors",
385
+ "vision_backbone.featurizer.blocks.12.ls1.scale_factor": "model-00001-of-00004.safetensors",
386
+ "vision_backbone.featurizer.blocks.12.ls2.scale_factor": "model-00001-of-00004.safetensors",
387
+ "vision_backbone.featurizer.blocks.12.mlp.fc1.bias": "model-00001-of-00004.safetensors",
388
+ "vision_backbone.featurizer.blocks.12.mlp.fc1.weight": "model-00001-of-00004.safetensors",
389
+ "vision_backbone.featurizer.blocks.12.mlp.fc2.bias": "model-00001-of-00004.safetensors",
390
+ "vision_backbone.featurizer.blocks.12.mlp.fc2.weight": "model-00001-of-00004.safetensors",
391
+ "vision_backbone.featurizer.blocks.12.norm1.bias": "model-00001-of-00004.safetensors",
392
+ "vision_backbone.featurizer.blocks.12.norm1.weight": "model-00001-of-00004.safetensors",
393
+ "vision_backbone.featurizer.blocks.12.norm2.bias": "model-00001-of-00004.safetensors",
394
+ "vision_backbone.featurizer.blocks.12.norm2.weight": "model-00001-of-00004.safetensors",
395
+ "vision_backbone.featurizer.blocks.13.attn.proj.bias": "model-00001-of-00004.safetensors",
396
+ "vision_backbone.featurizer.blocks.13.attn.proj.weight": "model-00001-of-00004.safetensors",
397
+ "vision_backbone.featurizer.blocks.13.attn.qkv.bias": "model-00001-of-00004.safetensors",
398
+ "vision_backbone.featurizer.blocks.13.attn.qkv.weight": "model-00001-of-00004.safetensors",
399
+ "vision_backbone.featurizer.blocks.13.ls1.scale_factor": "model-00001-of-00004.safetensors",
400
+ "vision_backbone.featurizer.blocks.13.ls2.scale_factor": "model-00001-of-00004.safetensors",
401
+ "vision_backbone.featurizer.blocks.13.mlp.fc1.bias": "model-00001-of-00004.safetensors",
402
+ "vision_backbone.featurizer.blocks.13.mlp.fc1.weight": "model-00001-of-00004.safetensors",
403
+ "vision_backbone.featurizer.blocks.13.mlp.fc2.bias": "model-00001-of-00004.safetensors",
404
+ "vision_backbone.featurizer.blocks.13.mlp.fc2.weight": "model-00001-of-00004.safetensors",
405
+ "vision_backbone.featurizer.blocks.13.norm1.bias": "model-00001-of-00004.safetensors",
406
+ "vision_backbone.featurizer.blocks.13.norm1.weight": "model-00001-of-00004.safetensors",
407
+ "vision_backbone.featurizer.blocks.13.norm2.bias": "model-00001-of-00004.safetensors",
408
+ "vision_backbone.featurizer.blocks.13.norm2.weight": "model-00001-of-00004.safetensors",
409
+ "vision_backbone.featurizer.blocks.14.attn.proj.bias": "model-00001-of-00004.safetensors",
410
+ "vision_backbone.featurizer.blocks.14.attn.proj.weight": "model-00001-of-00004.safetensors",
411
+ "vision_backbone.featurizer.blocks.14.attn.qkv.bias": "model-00001-of-00004.safetensors",
412
+ "vision_backbone.featurizer.blocks.14.attn.qkv.weight": "model-00001-of-00004.safetensors",
413
+ "vision_backbone.featurizer.blocks.14.ls1.scale_factor": "model-00001-of-00004.safetensors",
414
+ "vision_backbone.featurizer.blocks.14.ls2.scale_factor": "model-00001-of-00004.safetensors",
415
+ "vision_backbone.featurizer.blocks.14.mlp.fc1.bias": "model-00001-of-00004.safetensors",
416
+ "vision_backbone.featurizer.blocks.14.mlp.fc1.weight": "model-00001-of-00004.safetensors",
417
+ "vision_backbone.featurizer.blocks.14.mlp.fc2.bias": "model-00001-of-00004.safetensors",
418
+ "vision_backbone.featurizer.blocks.14.mlp.fc2.weight": "model-00001-of-00004.safetensors",
419
+ "vision_backbone.featurizer.blocks.14.norm1.bias": "model-00001-of-00004.safetensors",
420
+ "vision_backbone.featurizer.blocks.14.norm1.weight": "model-00001-of-00004.safetensors",
421
+ "vision_backbone.featurizer.blocks.14.norm2.bias": "model-00001-of-00004.safetensors",
422
+ "vision_backbone.featurizer.blocks.14.norm2.weight": "model-00001-of-00004.safetensors",
423
+ "vision_backbone.featurizer.blocks.15.attn.proj.bias": "model-00001-of-00004.safetensors",
424
+ "vision_backbone.featurizer.blocks.15.attn.proj.weight": "model-00001-of-00004.safetensors",
425
+ "vision_backbone.featurizer.blocks.15.attn.qkv.bias": "model-00001-of-00004.safetensors",
426
+ "vision_backbone.featurizer.blocks.15.attn.qkv.weight": "model-00001-of-00004.safetensors",
427
+ "vision_backbone.featurizer.blocks.15.ls1.scale_factor": "model-00001-of-00004.safetensors",
428
+ "vision_backbone.featurizer.blocks.15.ls2.scale_factor": "model-00001-of-00004.safetensors",
429
+ "vision_backbone.featurizer.blocks.15.mlp.fc1.bias": "model-00001-of-00004.safetensors",
430
+ "vision_backbone.featurizer.blocks.15.mlp.fc1.weight": "model-00001-of-00004.safetensors",
431
+ "vision_backbone.featurizer.blocks.15.mlp.fc2.bias": "model-00001-of-00004.safetensors",
432
+ "vision_backbone.featurizer.blocks.15.mlp.fc2.weight": "model-00001-of-00004.safetensors",
433
+ "vision_backbone.featurizer.blocks.15.norm1.bias": "model-00001-of-00004.safetensors",
434
+ "vision_backbone.featurizer.blocks.15.norm1.weight": "model-00001-of-00004.safetensors",
435
+ "vision_backbone.featurizer.blocks.15.norm2.bias": "model-00001-of-00004.safetensors",
436
+ "vision_backbone.featurizer.blocks.15.norm2.weight": "model-00001-of-00004.safetensors",
437
+ "vision_backbone.featurizer.blocks.16.attn.proj.bias": "model-00001-of-00004.safetensors",
438
+ "vision_backbone.featurizer.blocks.16.attn.proj.weight": "model-00001-of-00004.safetensors",
439
+ "vision_backbone.featurizer.blocks.16.attn.qkv.bias": "model-00001-of-00004.safetensors",
440
+ "vision_backbone.featurizer.blocks.16.attn.qkv.weight": "model-00001-of-00004.safetensors",
441
+ "vision_backbone.featurizer.blocks.16.ls1.scale_factor": "model-00001-of-00004.safetensors",
442
+ "vision_backbone.featurizer.blocks.16.ls2.scale_factor": "model-00001-of-00004.safetensors",
443
+ "vision_backbone.featurizer.blocks.16.mlp.fc1.bias": "model-00001-of-00004.safetensors",
444
+ "vision_backbone.featurizer.blocks.16.mlp.fc1.weight": "model-00001-of-00004.safetensors",
445
+ "vision_backbone.featurizer.blocks.16.mlp.fc2.bias": "model-00001-of-00004.safetensors",
446
+ "vision_backbone.featurizer.blocks.16.mlp.fc2.weight": "model-00001-of-00004.safetensors",
447
+ "vision_backbone.featurizer.blocks.16.norm1.bias": "model-00001-of-00004.safetensors",
448
+ "vision_backbone.featurizer.blocks.16.norm1.weight": "model-00001-of-00004.safetensors",
449
+ "vision_backbone.featurizer.blocks.16.norm2.bias": "model-00001-of-00004.safetensors",
450
+ "vision_backbone.featurizer.blocks.16.norm2.weight": "model-00001-of-00004.safetensors",
451
+ "vision_backbone.featurizer.blocks.17.attn.proj.bias": "model-00001-of-00004.safetensors",
452
+ "vision_backbone.featurizer.blocks.17.attn.proj.weight": "model-00001-of-00004.safetensors",
453
+ "vision_backbone.featurizer.blocks.17.attn.qkv.bias": "model-00001-of-00004.safetensors",
454
+ "vision_backbone.featurizer.blocks.17.attn.qkv.weight": "model-00001-of-00004.safetensors",
455
+ "vision_backbone.featurizer.blocks.17.ls1.scale_factor": "model-00001-of-00004.safetensors",
456
+ "vision_backbone.featurizer.blocks.17.ls2.scale_factor": "model-00001-of-00004.safetensors",
457
+ "vision_backbone.featurizer.blocks.17.mlp.fc1.bias": "model-00001-of-00004.safetensors",
458
+ "vision_backbone.featurizer.blocks.17.mlp.fc1.weight": "model-00001-of-00004.safetensors",
459
+ "vision_backbone.featurizer.blocks.17.mlp.fc2.bias": "model-00001-of-00004.safetensors",
460
+ "vision_backbone.featurizer.blocks.17.mlp.fc2.weight": "model-00001-of-00004.safetensors",
461
+ "vision_backbone.featurizer.blocks.17.norm1.bias": "model-00001-of-00004.safetensors",
462
+ "vision_backbone.featurizer.blocks.17.norm1.weight": "model-00001-of-00004.safetensors",
463
+ "vision_backbone.featurizer.blocks.17.norm2.bias": "model-00001-of-00004.safetensors",
464
+ "vision_backbone.featurizer.blocks.17.norm2.weight": "model-00001-of-00004.safetensors",
465
+ "vision_backbone.featurizer.blocks.18.attn.proj.bias": "model-00001-of-00004.safetensors",
466
+ "vision_backbone.featurizer.blocks.18.attn.proj.weight": "model-00001-of-00004.safetensors",
467
+ "vision_backbone.featurizer.blocks.18.attn.qkv.bias": "model-00001-of-00004.safetensors",
468
+ "vision_backbone.featurizer.blocks.18.attn.qkv.weight": "model-00001-of-00004.safetensors",
469
+ "vision_backbone.featurizer.blocks.18.ls1.scale_factor": "model-00001-of-00004.safetensors",
470
+ "vision_backbone.featurizer.blocks.18.ls2.scale_factor": "model-00001-of-00004.safetensors",
471
+ "vision_backbone.featurizer.blocks.18.mlp.fc1.bias": "model-00001-of-00004.safetensors",
472
+ "vision_backbone.featurizer.blocks.18.mlp.fc1.weight": "model-00001-of-00004.safetensors",
473
+ "vision_backbone.featurizer.blocks.18.mlp.fc2.bias": "model-00001-of-00004.safetensors",
474
+ "vision_backbone.featurizer.blocks.18.mlp.fc2.weight": "model-00001-of-00004.safetensors",
475
+ "vision_backbone.featurizer.blocks.18.norm1.bias": "model-00001-of-00004.safetensors",
476
+ "vision_backbone.featurizer.blocks.18.norm1.weight": "model-00001-of-00004.safetensors",
477
+ "vision_backbone.featurizer.blocks.18.norm2.bias": "model-00001-of-00004.safetensors",
478
+ "vision_backbone.featurizer.blocks.18.norm2.weight": "model-00001-of-00004.safetensors",
479
+ "vision_backbone.featurizer.blocks.19.attn.proj.bias": "model-00001-of-00004.safetensors",
480
+ "vision_backbone.featurizer.blocks.19.attn.proj.weight": "model-00001-of-00004.safetensors",
481
+ "vision_backbone.featurizer.blocks.19.attn.qkv.bias": "model-00001-of-00004.safetensors",
482
+ "vision_backbone.featurizer.blocks.19.attn.qkv.weight": "model-00001-of-00004.safetensors",
483
+ "vision_backbone.featurizer.blocks.19.ls1.scale_factor": "model-00001-of-00004.safetensors",
484
+ "vision_backbone.featurizer.blocks.19.ls2.scale_factor": "model-00001-of-00004.safetensors",
485
+ "vision_backbone.featurizer.blocks.19.mlp.fc1.bias": "model-00001-of-00004.safetensors",
486
+ "vision_backbone.featurizer.blocks.19.mlp.fc1.weight": "model-00001-of-00004.safetensors",
487
+ "vision_backbone.featurizer.blocks.19.mlp.fc2.bias": "model-00001-of-00004.safetensors",
488
+ "vision_backbone.featurizer.blocks.19.mlp.fc2.weight": "model-00001-of-00004.safetensors",
489
+ "vision_backbone.featurizer.blocks.19.norm1.bias": "model-00001-of-00004.safetensors",
490
+ "vision_backbone.featurizer.blocks.19.norm1.weight": "model-00001-of-00004.safetensors",
491
+ "vision_backbone.featurizer.blocks.19.norm2.bias": "model-00001-of-00004.safetensors",
492
+ "vision_backbone.featurizer.blocks.19.norm2.weight": "model-00001-of-00004.safetensors",
493
+ "vision_backbone.featurizer.blocks.2.attn.proj.bias": "model-00001-of-00004.safetensors",
494
+ "vision_backbone.featurizer.blocks.2.attn.proj.weight": "model-00001-of-00004.safetensors",
495
+ "vision_backbone.featurizer.blocks.2.attn.qkv.bias": "model-00001-of-00004.safetensors",
496
+ "vision_backbone.featurizer.blocks.2.attn.qkv.weight": "model-00001-of-00004.safetensors",
497
+ "vision_backbone.featurizer.blocks.2.ls1.scale_factor": "model-00001-of-00004.safetensors",
498
+ "vision_backbone.featurizer.blocks.2.ls2.scale_factor": "model-00001-of-00004.safetensors",
499
+ "vision_backbone.featurizer.blocks.2.mlp.fc1.bias": "model-00001-of-00004.safetensors",
500
+ "vision_backbone.featurizer.blocks.2.mlp.fc1.weight": "model-00001-of-00004.safetensors",
501
+ "vision_backbone.featurizer.blocks.2.mlp.fc2.bias": "model-00001-of-00004.safetensors",
502
+ "vision_backbone.featurizer.blocks.2.mlp.fc2.weight": "model-00001-of-00004.safetensors",
503
+ "vision_backbone.featurizer.blocks.2.norm1.bias": "model-00001-of-00004.safetensors",
504
+ "vision_backbone.featurizer.blocks.2.norm1.weight": "model-00001-of-00004.safetensors",
505
+ "vision_backbone.featurizer.blocks.2.norm2.bias": "model-00001-of-00004.safetensors",
506
+ "vision_backbone.featurizer.blocks.2.norm2.weight": "model-00001-of-00004.safetensors",
507
+ "vision_backbone.featurizer.blocks.20.attn.proj.bias": "model-00001-of-00004.safetensors",
508
+ "vision_backbone.featurizer.blocks.20.attn.proj.weight": "model-00001-of-00004.safetensors",
509
+ "vision_backbone.featurizer.blocks.20.attn.qkv.bias": "model-00001-of-00004.safetensors",
510
+ "vision_backbone.featurizer.blocks.20.attn.qkv.weight": "model-00001-of-00004.safetensors",
511
+ "vision_backbone.featurizer.blocks.20.ls1.scale_factor": "model-00001-of-00004.safetensors",
512
+ "vision_backbone.featurizer.blocks.20.ls2.scale_factor": "model-00001-of-00004.safetensors",
513
+ "vision_backbone.featurizer.blocks.20.mlp.fc1.bias": "model-00001-of-00004.safetensors",
514
+ "vision_backbone.featurizer.blocks.20.mlp.fc1.weight": "model-00001-of-00004.safetensors",
515
+ "vision_backbone.featurizer.blocks.20.mlp.fc2.bias": "model-00001-of-00004.safetensors",
516
+ "vision_backbone.featurizer.blocks.20.mlp.fc2.weight": "model-00001-of-00004.safetensors",
517
+ "vision_backbone.featurizer.blocks.20.norm1.bias": "model-00001-of-00004.safetensors",
518
+ "vision_backbone.featurizer.blocks.20.norm1.weight": "model-00001-of-00004.safetensors",
519
+ "vision_backbone.featurizer.blocks.20.norm2.bias": "model-00001-of-00004.safetensors",
520
+ "vision_backbone.featurizer.blocks.20.norm2.weight": "model-00001-of-00004.safetensors",
521
+ "vision_backbone.featurizer.blocks.21.attn.proj.bias": "model-00001-of-00004.safetensors",
522
+ "vision_backbone.featurizer.blocks.21.attn.proj.weight": "model-00001-of-00004.safetensors",
523
+ "vision_backbone.featurizer.blocks.21.attn.qkv.bias": "model-00001-of-00004.safetensors",
524
+ "vision_backbone.featurizer.blocks.21.attn.qkv.weight": "model-00001-of-00004.safetensors",
525
+ "vision_backbone.featurizer.blocks.21.ls1.scale_factor": "model-00001-of-00004.safetensors",
526
+ "vision_backbone.featurizer.blocks.21.ls2.scale_factor": "model-00001-of-00004.safetensors",
527
+ "vision_backbone.featurizer.blocks.21.mlp.fc1.bias": "model-00001-of-00004.safetensors",
528
+ "vision_backbone.featurizer.blocks.21.mlp.fc1.weight": "model-00001-of-00004.safetensors",
529
+ "vision_backbone.featurizer.blocks.21.mlp.fc2.bias": "model-00001-of-00004.safetensors",
530
+ "vision_backbone.featurizer.blocks.21.mlp.fc2.weight": "model-00001-of-00004.safetensors",
531
+ "vision_backbone.featurizer.blocks.21.norm1.bias": "model-00001-of-00004.safetensors",
532
+ "vision_backbone.featurizer.blocks.21.norm1.weight": "model-00001-of-00004.safetensors",
533
+ "vision_backbone.featurizer.blocks.21.norm2.bias": "model-00001-of-00004.safetensors",
534
+ "vision_backbone.featurizer.blocks.21.norm2.weight": "model-00001-of-00004.safetensors",
535
+ "vision_backbone.featurizer.blocks.22.attn.proj.bias": "model-00001-of-00004.safetensors",
536
+ "vision_backbone.featurizer.blocks.22.attn.proj.weight": "model-00001-of-00004.safetensors",
537
+ "vision_backbone.featurizer.blocks.22.attn.qkv.bias": "model-00001-of-00004.safetensors",
538
+ "vision_backbone.featurizer.blocks.22.attn.qkv.weight": "model-00001-of-00004.safetensors",
539
+ "vision_backbone.featurizer.blocks.22.ls1.scale_factor": "model-00001-of-00004.safetensors",
540
+ "vision_backbone.featurizer.blocks.22.ls2.scale_factor": "model-00001-of-00004.safetensors",
541
+ "vision_backbone.featurizer.blocks.22.mlp.fc1.bias": "model-00001-of-00004.safetensors",
542
+ "vision_backbone.featurizer.blocks.22.mlp.fc1.weight": "model-00001-of-00004.safetensors",
543
+ "vision_backbone.featurizer.blocks.22.mlp.fc2.bias": "model-00001-of-00004.safetensors",
544
+ "vision_backbone.featurizer.blocks.22.mlp.fc2.weight": "model-00001-of-00004.safetensors",
545
+ "vision_backbone.featurizer.blocks.22.norm1.bias": "model-00001-of-00004.safetensors",
546
+ "vision_backbone.featurizer.blocks.22.norm1.weight": "model-00001-of-00004.safetensors",
547
+ "vision_backbone.featurizer.blocks.22.norm2.bias": "model-00001-of-00004.safetensors",
548
+ "vision_backbone.featurizer.blocks.22.norm2.weight": "model-00001-of-00004.safetensors",
549
+ "vision_backbone.featurizer.blocks.23.attn.proj.bias": "model-00001-of-00004.safetensors",
550
+ "vision_backbone.featurizer.blocks.23.attn.proj.weight": "model-00001-of-00004.safetensors",
551
+ "vision_backbone.featurizer.blocks.23.attn.qkv.bias": "model-00001-of-00004.safetensors",
552
+ "vision_backbone.featurizer.blocks.23.attn.qkv.weight": "model-00001-of-00004.safetensors",
553
+ "vision_backbone.featurizer.blocks.23.ls1.scale_factor": "model-00001-of-00004.safetensors",
554
+ "vision_backbone.featurizer.blocks.23.ls2.scale_factor": "model-00001-of-00004.safetensors",
555
+ "vision_backbone.featurizer.blocks.23.mlp.fc1.bias": "model-00001-of-00004.safetensors",
556
+ "vision_backbone.featurizer.blocks.23.mlp.fc1.weight": "model-00001-of-00004.safetensors",
557
+ "vision_backbone.featurizer.blocks.23.mlp.fc2.bias": "model-00001-of-00004.safetensors",
558
+ "vision_backbone.featurizer.blocks.23.mlp.fc2.weight": "model-00001-of-00004.safetensors",
559
+ "vision_backbone.featurizer.blocks.23.norm1.bias": "model-00001-of-00004.safetensors",
560
+ "vision_backbone.featurizer.blocks.23.norm1.weight": "model-00001-of-00004.safetensors",
561
+ "vision_backbone.featurizer.blocks.23.norm2.bias": "model-00001-of-00004.safetensors",
562
+ "vision_backbone.featurizer.blocks.23.norm2.weight": "model-00001-of-00004.safetensors",
563
+ "vision_backbone.featurizer.blocks.3.attn.proj.bias": "model-00001-of-00004.safetensors",
564
+ "vision_backbone.featurizer.blocks.3.attn.proj.weight": "model-00001-of-00004.safetensors",
565
+ "vision_backbone.featurizer.blocks.3.attn.qkv.bias": "model-00001-of-00004.safetensors",
566
+ "vision_backbone.featurizer.blocks.3.attn.qkv.weight": "model-00001-of-00004.safetensors",
567
+ "vision_backbone.featurizer.blocks.3.ls1.scale_factor": "model-00001-of-00004.safetensors",
568
+ "vision_backbone.featurizer.blocks.3.ls2.scale_factor": "model-00001-of-00004.safetensors",
569
+ "vision_backbone.featurizer.blocks.3.mlp.fc1.bias": "model-00001-of-00004.safetensors",
570
+ "vision_backbone.featurizer.blocks.3.mlp.fc1.weight": "model-00001-of-00004.safetensors",
571
+ "vision_backbone.featurizer.blocks.3.mlp.fc2.bias": "model-00001-of-00004.safetensors",
572
+ "vision_backbone.featurizer.blocks.3.mlp.fc2.weight": "model-00001-of-00004.safetensors",
573
+ "vision_backbone.featurizer.blocks.3.norm1.bias": "model-00001-of-00004.safetensors",
574
+ "vision_backbone.featurizer.blocks.3.norm1.weight": "model-00001-of-00004.safetensors",
575
+ "vision_backbone.featurizer.blocks.3.norm2.bias": "model-00001-of-00004.safetensors",
576
+ "vision_backbone.featurizer.blocks.3.norm2.weight": "model-00001-of-00004.safetensors",
577
+ "vision_backbone.featurizer.blocks.4.attn.proj.bias": "model-00001-of-00004.safetensors",
578
+ "vision_backbone.featurizer.blocks.4.attn.proj.weight": "model-00001-of-00004.safetensors",
579
+ "vision_backbone.featurizer.blocks.4.attn.qkv.bias": "model-00001-of-00004.safetensors",
580
+ "vision_backbone.featurizer.blocks.4.attn.qkv.weight": "model-00001-of-00004.safetensors",
581
+ "vision_backbone.featurizer.blocks.4.ls1.scale_factor": "model-00001-of-00004.safetensors",
582
+ "vision_backbone.featurizer.blocks.4.ls2.scale_factor": "model-00001-of-00004.safetensors",
583
+ "vision_backbone.featurizer.blocks.4.mlp.fc1.bias": "model-00001-of-00004.safetensors",
584
+ "vision_backbone.featurizer.blocks.4.mlp.fc1.weight": "model-00001-of-00004.safetensors",
585
+ "vision_backbone.featurizer.blocks.4.mlp.fc2.bias": "model-00001-of-00004.safetensors",
586
+ "vision_backbone.featurizer.blocks.4.mlp.fc2.weight": "model-00001-of-00004.safetensors",
587
+ "vision_backbone.featurizer.blocks.4.norm1.bias": "model-00001-of-00004.safetensors",
588
+ "vision_backbone.featurizer.blocks.4.norm1.weight": "model-00001-of-00004.safetensors",
589
+ "vision_backbone.featurizer.blocks.4.norm2.bias": "model-00001-of-00004.safetensors",
590
+ "vision_backbone.featurizer.blocks.4.norm2.weight": "model-00001-of-00004.safetensors",
591
+ "vision_backbone.featurizer.blocks.5.attn.proj.bias": "model-00001-of-00004.safetensors",
592
+ "vision_backbone.featurizer.blocks.5.attn.proj.weight": "model-00001-of-00004.safetensors",
593
+ "vision_backbone.featurizer.blocks.5.attn.qkv.bias": "model-00001-of-00004.safetensors",
594
+ "vision_backbone.featurizer.blocks.5.attn.qkv.weight": "model-00001-of-00004.safetensors",
595
+ "vision_backbone.featurizer.blocks.5.ls1.scale_factor": "model-00001-of-00004.safetensors",
596
+ "vision_backbone.featurizer.blocks.5.ls2.scale_factor": "model-00001-of-00004.safetensors",
597
+ "vision_backbone.featurizer.blocks.5.mlp.fc1.bias": "model-00001-of-00004.safetensors",
598
+ "vision_backbone.featurizer.blocks.5.mlp.fc1.weight": "model-00001-of-00004.safetensors",
599
+ "vision_backbone.featurizer.blocks.5.mlp.fc2.bias": "model-00001-of-00004.safetensors",
600
+ "vision_backbone.featurizer.blocks.5.mlp.fc2.weight": "model-00001-of-00004.safetensors",
601
+ "vision_backbone.featurizer.blocks.5.norm1.bias": "model-00001-of-00004.safetensors",
602
+ "vision_backbone.featurizer.blocks.5.norm1.weight": "model-00001-of-00004.safetensors",
603
+ "vision_backbone.featurizer.blocks.5.norm2.bias": "model-00001-of-00004.safetensors",
604
+ "vision_backbone.featurizer.blocks.5.norm2.weight": "model-00001-of-00004.safetensors",
605
+ "vision_backbone.featurizer.blocks.6.attn.proj.bias": "model-00001-of-00004.safetensors",
606
+ "vision_backbone.featurizer.blocks.6.attn.proj.weight": "model-00001-of-00004.safetensors",
607
+ "vision_backbone.featurizer.blocks.6.attn.qkv.bias": "model-00001-of-00004.safetensors",
608
+ "vision_backbone.featurizer.blocks.6.attn.qkv.weight": "model-00001-of-00004.safetensors",
609
+ "vision_backbone.featurizer.blocks.6.ls1.scale_factor": "model-00001-of-00004.safetensors",
610
+ "vision_backbone.featurizer.blocks.6.ls2.scale_factor": "model-00001-of-00004.safetensors",
611
+ "vision_backbone.featurizer.blocks.6.mlp.fc1.bias": "model-00001-of-00004.safetensors",
612
+ "vision_backbone.featurizer.blocks.6.mlp.fc1.weight": "model-00001-of-00004.safetensors",
613
+ "vision_backbone.featurizer.blocks.6.mlp.fc2.bias": "model-00001-of-00004.safetensors",
614
+ "vision_backbone.featurizer.blocks.6.mlp.fc2.weight": "model-00001-of-00004.safetensors",
615
+ "vision_backbone.featurizer.blocks.6.norm1.bias": "model-00001-of-00004.safetensors",
616
+ "vision_backbone.featurizer.blocks.6.norm1.weight": "model-00001-of-00004.safetensors",
617
+ "vision_backbone.featurizer.blocks.6.norm2.bias": "model-00001-of-00004.safetensors",
618
+ "vision_backbone.featurizer.blocks.6.norm2.weight": "model-00001-of-00004.safetensors",
619
+ "vision_backbone.featurizer.blocks.7.attn.proj.bias": "model-00001-of-00004.safetensors",
620
+ "vision_backbone.featurizer.blocks.7.attn.proj.weight": "model-00001-of-00004.safetensors",
621
+ "vision_backbone.featurizer.blocks.7.attn.qkv.bias": "model-00001-of-00004.safetensors",
622
+ "vision_backbone.featurizer.blocks.7.attn.qkv.weight": "model-00001-of-00004.safetensors",
623
+ "vision_backbone.featurizer.blocks.7.ls1.scale_factor": "model-00001-of-00004.safetensors",
624
+ "vision_backbone.featurizer.blocks.7.ls2.scale_factor": "model-00001-of-00004.safetensors",
625
+ "vision_backbone.featurizer.blocks.7.mlp.fc1.bias": "model-00001-of-00004.safetensors",
626
+ "vision_backbone.featurizer.blocks.7.mlp.fc1.weight": "model-00001-of-00004.safetensors",
627
+ "vision_backbone.featurizer.blocks.7.mlp.fc2.bias": "model-00001-of-00004.safetensors",
628
+ "vision_backbone.featurizer.blocks.7.mlp.fc2.weight": "model-00001-of-00004.safetensors",
629
+ "vision_backbone.featurizer.blocks.7.norm1.bias": "model-00001-of-00004.safetensors",
630
+ "vision_backbone.featurizer.blocks.7.norm1.weight": "model-00001-of-00004.safetensors",
631
+ "vision_backbone.featurizer.blocks.7.norm2.bias": "model-00001-of-00004.safetensors",
632
+ "vision_backbone.featurizer.blocks.7.norm2.weight": "model-00001-of-00004.safetensors",
633
+ "vision_backbone.featurizer.blocks.8.attn.proj.bias": "model-00001-of-00004.safetensors",
634
+ "vision_backbone.featurizer.blocks.8.attn.proj.weight": "model-00001-of-00004.safetensors",
635
+ "vision_backbone.featurizer.blocks.8.attn.qkv.bias": "model-00001-of-00004.safetensors",
636
+ "vision_backbone.featurizer.blocks.8.attn.qkv.weight": "model-00001-of-00004.safetensors",
637
+ "vision_backbone.featurizer.blocks.8.ls1.scale_factor": "model-00001-of-00004.safetensors",
638
+ "vision_backbone.featurizer.blocks.8.ls2.scale_factor": "model-00001-of-00004.safetensors",
639
+ "vision_backbone.featurizer.blocks.8.mlp.fc1.bias": "model-00001-of-00004.safetensors",
640
+ "vision_backbone.featurizer.blocks.8.mlp.fc1.weight": "model-00001-of-00004.safetensors",
641
+ "vision_backbone.featurizer.blocks.8.mlp.fc2.bias": "model-00001-of-00004.safetensors",
642
+ "vision_backbone.featurizer.blocks.8.mlp.fc2.weight": "model-00001-of-00004.safetensors",
643
+ "vision_backbone.featurizer.blocks.8.norm1.bias": "model-00001-of-00004.safetensors",
644
+ "vision_backbone.featurizer.blocks.8.norm1.weight": "model-00001-of-00004.safetensors",
645
+ "vision_backbone.featurizer.blocks.8.norm2.bias": "model-00001-of-00004.safetensors",
646
+ "vision_backbone.featurizer.blocks.8.norm2.weight": "model-00001-of-00004.safetensors",
647
+ "vision_backbone.featurizer.blocks.9.attn.proj.bias": "model-00001-of-00004.safetensors",
648
+ "vision_backbone.featurizer.blocks.9.attn.proj.weight": "model-00001-of-00004.safetensors",
649
+ "vision_backbone.featurizer.blocks.9.attn.qkv.bias": "model-00001-of-00004.safetensors",
650
+ "vision_backbone.featurizer.blocks.9.attn.qkv.weight": "model-00001-of-00004.safetensors",
651
+ "vision_backbone.featurizer.blocks.9.ls1.scale_factor": "model-00001-of-00004.safetensors",
652
+ "vision_backbone.featurizer.blocks.9.ls2.scale_factor": "model-00001-of-00004.safetensors",
653
+ "vision_backbone.featurizer.blocks.9.mlp.fc1.bias": "model-00001-of-00004.safetensors",
654
+ "vision_backbone.featurizer.blocks.9.mlp.fc1.weight": "model-00001-of-00004.safetensors",
655
+ "vision_backbone.featurizer.blocks.9.mlp.fc2.bias": "model-00001-of-00004.safetensors",
656
+ "vision_backbone.featurizer.blocks.9.mlp.fc2.weight": "model-00001-of-00004.safetensors",
657
+ "vision_backbone.featurizer.blocks.9.norm1.bias": "model-00001-of-00004.safetensors",
658
+ "vision_backbone.featurizer.blocks.9.norm1.weight": "model-00001-of-00004.safetensors",
659
+ "vision_backbone.featurizer.blocks.9.norm2.bias": "model-00001-of-00004.safetensors",
660
+ "vision_backbone.featurizer.blocks.9.norm2.weight": "model-00001-of-00004.safetensors",
661
+ "vision_backbone.featurizer.cls_token": "model-00001-of-00004.safetensors",
662
+ "vision_backbone.featurizer.norm.bias": "model-00001-of-00004.safetensors",
663
+ "vision_backbone.featurizer.norm.weight": "model-00001-of-00004.safetensors",
664
+ "vision_backbone.featurizer.patch_embed.proj.bias": "model-00001-of-00004.safetensors",
665
+ "vision_backbone.featurizer.patch_embed.proj.weight": "model-00001-of-00004.safetensors",
666
+ "vision_backbone.featurizer.pos_embed": "model-00001-of-00004.safetensors",
667
+ "vision_backbone.featurizer.reg_token": "model-00001-of-00004.safetensors",
668
+ "vision_backbone.fused_featurizer.attn_pool.kv.bias": "model-00001-of-00004.safetensors",
669
+ "vision_backbone.fused_featurizer.attn_pool.kv.weight": "model-00001-of-00004.safetensors",
670
+ "vision_backbone.fused_featurizer.attn_pool.latent": "model-00001-of-00004.safetensors",
671
+ "vision_backbone.fused_featurizer.attn_pool.mlp.fc1.bias": "model-00001-of-00004.safetensors",
672
+ "vision_backbone.fused_featurizer.attn_pool.mlp.fc1.weight": "model-00001-of-00004.safetensors",
673
+ "vision_backbone.fused_featurizer.attn_pool.mlp.fc2.bias": "model-00001-of-00004.safetensors",
674
+ "vision_backbone.fused_featurizer.attn_pool.mlp.fc2.weight": "model-00001-of-00004.safetensors",
675
+ "vision_backbone.fused_featurizer.attn_pool.norm.bias": "model-00001-of-00004.safetensors",
676
+ "vision_backbone.fused_featurizer.attn_pool.norm.weight": "model-00001-of-00004.safetensors",
677
+ "vision_backbone.fused_featurizer.attn_pool.proj.bias": "model-00001-of-00004.safetensors",
678
+ "vision_backbone.fused_featurizer.attn_pool.proj.weight": "model-00001-of-00004.safetensors",
679
+ "vision_backbone.fused_featurizer.attn_pool.q.bias": "model-00001-of-00004.safetensors",
680
+ "vision_backbone.fused_featurizer.attn_pool.q.weight": "model-00001-of-00004.safetensors",
681
+ "vision_backbone.fused_featurizer.blocks.0.attn.proj.bias": "model-00001-of-00004.safetensors",
682
+ "vision_backbone.fused_featurizer.blocks.0.attn.proj.weight": "model-00001-of-00004.safetensors",
683
+ "vision_backbone.fused_featurizer.blocks.0.attn.qkv.bias": "model-00001-of-00004.safetensors",
684
+ "vision_backbone.fused_featurizer.blocks.0.attn.qkv.weight": "model-00001-of-00004.safetensors",
685
+ "vision_backbone.fused_featurizer.blocks.0.mlp.fc1.bias": "model-00001-of-00004.safetensors",
686
+ "vision_backbone.fused_featurizer.blocks.0.mlp.fc1.weight": "model-00001-of-00004.safetensors",
687
+ "vision_backbone.fused_featurizer.blocks.0.mlp.fc2.bias": "model-00001-of-00004.safetensors",
688
+ "vision_backbone.fused_featurizer.blocks.0.mlp.fc2.weight": "model-00001-of-00004.safetensors",
689
+ "vision_backbone.fused_featurizer.blocks.0.norm1.bias": "model-00001-of-00004.safetensors",
690
+ "vision_backbone.fused_featurizer.blocks.0.norm1.weight": "model-00001-of-00004.safetensors",
691
+ "vision_backbone.fused_featurizer.blocks.0.norm2.bias": "model-00001-of-00004.safetensors",
692
+ "vision_backbone.fused_featurizer.blocks.0.norm2.weight": "model-00001-of-00004.safetensors",
693
+ "vision_backbone.fused_featurizer.blocks.1.attn.proj.bias": "model-00001-of-00004.safetensors",
694
+ "vision_backbone.fused_featurizer.blocks.1.attn.proj.weight": "model-00001-of-00004.safetensors",
695
+ "vision_backbone.fused_featurizer.blocks.1.attn.qkv.bias": "model-00001-of-00004.safetensors",
696
+ "vision_backbone.fused_featurizer.blocks.1.attn.qkv.weight": "model-00001-of-00004.safetensors",
697
+ "vision_backbone.fused_featurizer.blocks.1.mlp.fc1.bias": "model-00001-of-00004.safetensors",
698
+ "vision_backbone.fused_featurizer.blocks.1.mlp.fc1.weight": "model-00001-of-00004.safetensors",
699
+ "vision_backbone.fused_featurizer.blocks.1.mlp.fc2.bias": "model-00001-of-00004.safetensors",
700
+ "vision_backbone.fused_featurizer.blocks.1.mlp.fc2.weight": "model-00001-of-00004.safetensors",
701
+ "vision_backbone.fused_featurizer.blocks.1.norm1.bias": "model-00001-of-00004.safetensors",
702
+ "vision_backbone.fused_featurizer.blocks.1.norm1.weight": "model-00001-of-00004.safetensors",
703
+ "vision_backbone.fused_featurizer.blocks.1.norm2.bias": "model-00001-of-00004.safetensors",
704
+ "vision_backbone.fused_featurizer.blocks.1.norm2.weight": "model-00001-of-00004.safetensors",
705
+ "vision_backbone.fused_featurizer.blocks.10.attn.proj.bias": "model-00001-of-00004.safetensors",
706
+ "vision_backbone.fused_featurizer.blocks.10.attn.proj.weight": "model-00001-of-00004.safetensors",
707
+ "vision_backbone.fused_featurizer.blocks.10.attn.qkv.bias": "model-00001-of-00004.safetensors",
708
+ "vision_backbone.fused_featurizer.blocks.10.attn.qkv.weight": "model-00001-of-00004.safetensors",
709
+ "vision_backbone.fused_featurizer.blocks.10.mlp.fc1.bias": "model-00001-of-00004.safetensors",
710
+ "vision_backbone.fused_featurizer.blocks.10.mlp.fc1.weight": "model-00001-of-00004.safetensors",
711
+ "vision_backbone.fused_featurizer.blocks.10.mlp.fc2.bias": "model-00001-of-00004.safetensors",
712
+ "vision_backbone.fused_featurizer.blocks.10.mlp.fc2.weight": "model-00001-of-00004.safetensors",
713
+ "vision_backbone.fused_featurizer.blocks.10.norm1.bias": "model-00001-of-00004.safetensors",
714
+ "vision_backbone.fused_featurizer.blocks.10.norm1.weight": "model-00001-of-00004.safetensors",
715
+ "vision_backbone.fused_featurizer.blocks.10.norm2.bias": "model-00001-of-00004.safetensors",
716
+ "vision_backbone.fused_featurizer.blocks.10.norm2.weight": "model-00001-of-00004.safetensors",
717
+ "vision_backbone.fused_featurizer.blocks.11.attn.proj.bias": "model-00001-of-00004.safetensors",
718
+ "vision_backbone.fused_featurizer.blocks.11.attn.proj.weight": "model-00001-of-00004.safetensors",
719
+ "vision_backbone.fused_featurizer.blocks.11.attn.qkv.bias": "model-00001-of-00004.safetensors",
720
+ "vision_backbone.fused_featurizer.blocks.11.attn.qkv.weight": "model-00001-of-00004.safetensors",
721
+ "vision_backbone.fused_featurizer.blocks.11.mlp.fc1.bias": "model-00001-of-00004.safetensors",
722
+ "vision_backbone.fused_featurizer.blocks.11.mlp.fc1.weight": "model-00001-of-00004.safetensors",
723
+ "vision_backbone.fused_featurizer.blocks.11.mlp.fc2.bias": "model-00001-of-00004.safetensors",
724
+ "vision_backbone.fused_featurizer.blocks.11.mlp.fc2.weight": "model-00001-of-00004.safetensors",
725
+ "vision_backbone.fused_featurizer.blocks.11.norm1.bias": "model-00001-of-00004.safetensors",
726
+ "vision_backbone.fused_featurizer.blocks.11.norm1.weight": "model-00001-of-00004.safetensors",
727
+ "vision_backbone.fused_featurizer.blocks.11.norm2.bias": "model-00001-of-00004.safetensors",
728
+ "vision_backbone.fused_featurizer.blocks.11.norm2.weight": "model-00001-of-00004.safetensors",
729
+ "vision_backbone.fused_featurizer.blocks.12.attn.proj.bias": "model-00001-of-00004.safetensors",
730
+ "vision_backbone.fused_featurizer.blocks.12.attn.proj.weight": "model-00001-of-00004.safetensors",
731
+ "vision_backbone.fused_featurizer.blocks.12.attn.qkv.bias": "model-00001-of-00004.safetensors",
732
+ "vision_backbone.fused_featurizer.blocks.12.attn.qkv.weight": "model-00001-of-00004.safetensors",
733
+ "vision_backbone.fused_featurizer.blocks.12.mlp.fc1.bias": "model-00001-of-00004.safetensors",
734
+ "vision_backbone.fused_featurizer.blocks.12.mlp.fc1.weight": "model-00001-of-00004.safetensors",
735
+ "vision_backbone.fused_featurizer.blocks.12.mlp.fc2.bias": "model-00001-of-00004.safetensors",
736
+ "vision_backbone.fused_featurizer.blocks.12.mlp.fc2.weight": "model-00001-of-00004.safetensors",
737
+ "vision_backbone.fused_featurizer.blocks.12.norm1.bias": "model-00001-of-00004.safetensors",
738
+ "vision_backbone.fused_featurizer.blocks.12.norm1.weight": "model-00001-of-00004.safetensors",
739
+ "vision_backbone.fused_featurizer.blocks.12.norm2.bias": "model-00001-of-00004.safetensors",
740
+ "vision_backbone.fused_featurizer.blocks.12.norm2.weight": "model-00001-of-00004.safetensors",
741
+ "vision_backbone.fused_featurizer.blocks.13.attn.proj.bias": "model-00001-of-00004.safetensors",
742
+ "vision_backbone.fused_featurizer.blocks.13.attn.proj.weight": "model-00001-of-00004.safetensors",
743
+ "vision_backbone.fused_featurizer.blocks.13.attn.qkv.bias": "model-00001-of-00004.safetensors",
744
+ "vision_backbone.fused_featurizer.blocks.13.attn.qkv.weight": "model-00001-of-00004.safetensors",
745
+ "vision_backbone.fused_featurizer.blocks.13.mlp.fc1.bias": "model-00001-of-00004.safetensors",
746
+ "vision_backbone.fused_featurizer.blocks.13.mlp.fc1.weight": "model-00001-of-00004.safetensors",
747
+ "vision_backbone.fused_featurizer.blocks.13.mlp.fc2.bias": "model-00001-of-00004.safetensors",
748
+ "vision_backbone.fused_featurizer.blocks.13.mlp.fc2.weight": "model-00001-of-00004.safetensors",
749
+ "vision_backbone.fused_featurizer.blocks.13.norm1.bias": "model-00001-of-00004.safetensors",
750
+ "vision_backbone.fused_featurizer.blocks.13.norm1.weight": "model-00001-of-00004.safetensors",
751
+ "vision_backbone.fused_featurizer.blocks.13.norm2.bias": "model-00001-of-00004.safetensors",
752
+ "vision_backbone.fused_featurizer.blocks.13.norm2.weight": "model-00001-of-00004.safetensors",
753
+ "vision_backbone.fused_featurizer.blocks.14.attn.proj.bias": "model-00001-of-00004.safetensors",
754
+ "vision_backbone.fused_featurizer.blocks.14.attn.proj.weight": "model-00001-of-00004.safetensors",
755
+ "vision_backbone.fused_featurizer.blocks.14.attn.qkv.bias": "model-00001-of-00004.safetensors",
756
+ "vision_backbone.fused_featurizer.blocks.14.attn.qkv.weight": "model-00001-of-00004.safetensors",
757
+ "vision_backbone.fused_featurizer.blocks.14.mlp.fc1.bias": "model-00001-of-00004.safetensors",
758
+ "vision_backbone.fused_featurizer.blocks.14.mlp.fc1.weight": "model-00001-of-00004.safetensors",
759
+ "vision_backbone.fused_featurizer.blocks.14.mlp.fc2.bias": "model-00001-of-00004.safetensors",
760
+ "vision_backbone.fused_featurizer.blocks.14.mlp.fc2.weight": "model-00001-of-00004.safetensors",
761
+ "vision_backbone.fused_featurizer.blocks.14.norm1.bias": "model-00001-of-00004.safetensors",
762
+ "vision_backbone.fused_featurizer.blocks.14.norm1.weight": "model-00001-of-00004.safetensors",
763
+ "vision_backbone.fused_featurizer.blocks.14.norm2.bias": "model-00001-of-00004.safetensors",
764
+ "vision_backbone.fused_featurizer.blocks.14.norm2.weight": "model-00001-of-00004.safetensors",
765
+ "vision_backbone.fused_featurizer.blocks.15.attn.proj.bias": "model-00001-of-00004.safetensors",
766
+ "vision_backbone.fused_featurizer.blocks.15.attn.proj.weight": "model-00001-of-00004.safetensors",
767
+ "vision_backbone.fused_featurizer.blocks.15.attn.qkv.bias": "model-00001-of-00004.safetensors",
768
+ "vision_backbone.fused_featurizer.blocks.15.attn.qkv.weight": "model-00001-of-00004.safetensors",
769
+ "vision_backbone.fused_featurizer.blocks.15.mlp.fc1.bias": "model-00001-of-00004.safetensors",
770
+ "vision_backbone.fused_featurizer.blocks.15.mlp.fc1.weight": "model-00001-of-00004.safetensors",
771
+ "vision_backbone.fused_featurizer.blocks.15.mlp.fc2.bias": "model-00001-of-00004.safetensors",
772
+ "vision_backbone.fused_featurizer.blocks.15.mlp.fc2.weight": "model-00001-of-00004.safetensors",
773
+ "vision_backbone.fused_featurizer.blocks.15.norm1.bias": "model-00001-of-00004.safetensors",
774
+ "vision_backbone.fused_featurizer.blocks.15.norm1.weight": "model-00001-of-00004.safetensors",
775
+ "vision_backbone.fused_featurizer.blocks.15.norm2.bias": "model-00001-of-00004.safetensors",
776
+ "vision_backbone.fused_featurizer.blocks.15.norm2.weight": "model-00001-of-00004.safetensors",
777
+ "vision_backbone.fused_featurizer.blocks.16.attn.proj.bias": "model-00001-of-00004.safetensors",
778
+ "vision_backbone.fused_featurizer.blocks.16.attn.proj.weight": "model-00001-of-00004.safetensors",
779
+ "vision_backbone.fused_featurizer.blocks.16.attn.qkv.bias": "model-00001-of-00004.safetensors",
780
+ "vision_backbone.fused_featurizer.blocks.16.attn.qkv.weight": "model-00001-of-00004.safetensors",
781
+ "vision_backbone.fused_featurizer.blocks.16.mlp.fc1.bias": "model-00001-of-00004.safetensors",
782
+ "vision_backbone.fused_featurizer.blocks.16.mlp.fc1.weight": "model-00001-of-00004.safetensors",
783
+ "vision_backbone.fused_featurizer.blocks.16.mlp.fc2.bias": "model-00001-of-00004.safetensors",
784
+ "vision_backbone.fused_featurizer.blocks.16.mlp.fc2.weight": "model-00001-of-00004.safetensors",
785
+ "vision_backbone.fused_featurizer.blocks.16.norm1.bias": "model-00001-of-00004.safetensors",
786
+ "vision_backbone.fused_featurizer.blocks.16.norm1.weight": "model-00001-of-00004.safetensors",
787
+ "vision_backbone.fused_featurizer.blocks.16.norm2.bias": "model-00001-of-00004.safetensors",
788
+ "vision_backbone.fused_featurizer.blocks.16.norm2.weight": "model-00001-of-00004.safetensors",
789
+ "vision_backbone.fused_featurizer.blocks.17.attn.proj.bias": "model-00001-of-00004.safetensors",
790
+ "vision_backbone.fused_featurizer.blocks.17.attn.proj.weight": "model-00001-of-00004.safetensors",
791
+ "vision_backbone.fused_featurizer.blocks.17.attn.qkv.bias": "model-00001-of-00004.safetensors",
792
+ "vision_backbone.fused_featurizer.blocks.17.attn.qkv.weight": "model-00001-of-00004.safetensors",
793
+ "vision_backbone.fused_featurizer.blocks.17.mlp.fc1.bias": "model-00001-of-00004.safetensors",
794
+ "vision_backbone.fused_featurizer.blocks.17.mlp.fc1.weight": "model-00001-of-00004.safetensors",
795
+ "vision_backbone.fused_featurizer.blocks.17.mlp.fc2.bias": "model-00001-of-00004.safetensors",
796
+ "vision_backbone.fused_featurizer.blocks.17.mlp.fc2.weight": "model-00001-of-00004.safetensors",
797
+ "vision_backbone.fused_featurizer.blocks.17.norm1.bias": "model-00001-of-00004.safetensors",
798
+ "vision_backbone.fused_featurizer.blocks.17.norm1.weight": "model-00001-of-00004.safetensors",
799
+ "vision_backbone.fused_featurizer.blocks.17.norm2.bias": "model-00001-of-00004.safetensors",
800
+ "vision_backbone.fused_featurizer.blocks.17.norm2.weight": "model-00001-of-00004.safetensors",
801
+ "vision_backbone.fused_featurizer.blocks.18.attn.proj.bias": "model-00001-of-00004.safetensors",
802
+ "vision_backbone.fused_featurizer.blocks.18.attn.proj.weight": "model-00001-of-00004.safetensors",
803
+ "vision_backbone.fused_featurizer.blocks.18.attn.qkv.bias": "model-00001-of-00004.safetensors",
804
+ "vision_backbone.fused_featurizer.blocks.18.attn.qkv.weight": "model-00001-of-00004.safetensors",
805
+ "vision_backbone.fused_featurizer.blocks.18.mlp.fc1.bias": "model-00001-of-00004.safetensors",
806
+ "vision_backbone.fused_featurizer.blocks.18.mlp.fc1.weight": "model-00001-of-00004.safetensors",
807
+ "vision_backbone.fused_featurizer.blocks.18.mlp.fc2.bias": "model-00001-of-00004.safetensors",
808
+ "vision_backbone.fused_featurizer.blocks.18.mlp.fc2.weight": "model-00001-of-00004.safetensors",
809
+ "vision_backbone.fused_featurizer.blocks.18.norm1.bias": "model-00001-of-00004.safetensors",
810
+ "vision_backbone.fused_featurizer.blocks.18.norm1.weight": "model-00001-of-00004.safetensors",
811
+ "vision_backbone.fused_featurizer.blocks.18.norm2.bias": "model-00001-of-00004.safetensors",
812
+ "vision_backbone.fused_featurizer.blocks.18.norm2.weight": "model-00001-of-00004.safetensors",
813
+ "vision_backbone.fused_featurizer.blocks.19.attn.proj.bias": "model-00001-of-00004.safetensors",
814
+ "vision_backbone.fused_featurizer.blocks.19.attn.proj.weight": "model-00001-of-00004.safetensors",
815
+ "vision_backbone.fused_featurizer.blocks.19.attn.qkv.bias": "model-00001-of-00004.safetensors",
816
+ "vision_backbone.fused_featurizer.blocks.19.attn.qkv.weight": "model-00001-of-00004.safetensors",
817
+ "vision_backbone.fused_featurizer.blocks.19.mlp.fc1.bias": "model-00001-of-00004.safetensors",
818
+ "vision_backbone.fused_featurizer.blocks.19.mlp.fc1.weight": "model-00001-of-00004.safetensors",
819
+ "vision_backbone.fused_featurizer.blocks.19.mlp.fc2.bias": "model-00001-of-00004.safetensors",
820
+ "vision_backbone.fused_featurizer.blocks.19.mlp.fc2.weight": "model-00001-of-00004.safetensors",
821
+ "vision_backbone.fused_featurizer.blocks.19.norm1.bias": "model-00001-of-00004.safetensors",
822
+ "vision_backbone.fused_featurizer.blocks.19.norm1.weight": "model-00001-of-00004.safetensors",
823
+ "vision_backbone.fused_featurizer.blocks.19.norm2.bias": "model-00001-of-00004.safetensors",
824
+ "vision_backbone.fused_featurizer.blocks.19.norm2.weight": "model-00001-of-00004.safetensors",
825
+ "vision_backbone.fused_featurizer.blocks.2.attn.proj.bias": "model-00001-of-00004.safetensors",
826
+ "vision_backbone.fused_featurizer.blocks.2.attn.proj.weight": "model-00001-of-00004.safetensors",
827
+ "vision_backbone.fused_featurizer.blocks.2.attn.qkv.bias": "model-00001-of-00004.safetensors",
828
+ "vision_backbone.fused_featurizer.blocks.2.attn.qkv.weight": "model-00001-of-00004.safetensors",
829
+ "vision_backbone.fused_featurizer.blocks.2.mlp.fc1.bias": "model-00001-of-00004.safetensors",
830
+ "vision_backbone.fused_featurizer.blocks.2.mlp.fc1.weight": "model-00001-of-00004.safetensors",
831
+ "vision_backbone.fused_featurizer.blocks.2.mlp.fc2.bias": "model-00001-of-00004.safetensors",
832
+ "vision_backbone.fused_featurizer.blocks.2.mlp.fc2.weight": "model-00001-of-00004.safetensors",
833
+ "vision_backbone.fused_featurizer.blocks.2.norm1.bias": "model-00001-of-00004.safetensors",
834
+ "vision_backbone.fused_featurizer.blocks.2.norm1.weight": "model-00001-of-00004.safetensors",
835
+ "vision_backbone.fused_featurizer.blocks.2.norm2.bias": "model-00001-of-00004.safetensors",
836
+ "vision_backbone.fused_featurizer.blocks.2.norm2.weight": "model-00001-of-00004.safetensors",
837
+ "vision_backbone.fused_featurizer.blocks.20.attn.proj.bias": "model-00001-of-00004.safetensors",
838
+ "vision_backbone.fused_featurizer.blocks.20.attn.proj.weight": "model-00001-of-00004.safetensors",
839
+ "vision_backbone.fused_featurizer.blocks.20.attn.qkv.bias": "model-00001-of-00004.safetensors",
840
+ "vision_backbone.fused_featurizer.blocks.20.attn.qkv.weight": "model-00001-of-00004.safetensors",
841
+ "vision_backbone.fused_featurizer.blocks.20.mlp.fc1.bias": "model-00001-of-00004.safetensors",
842
+ "vision_backbone.fused_featurizer.blocks.20.mlp.fc1.weight": "model-00001-of-00004.safetensors",
843
+ "vision_backbone.fused_featurizer.blocks.20.mlp.fc2.bias": "model-00001-of-00004.safetensors",
844
+ "vision_backbone.fused_featurizer.blocks.20.mlp.fc2.weight": "model-00001-of-00004.safetensors",
845
+ "vision_backbone.fused_featurizer.blocks.20.norm1.bias": "model-00001-of-00004.safetensors",
846
+ "vision_backbone.fused_featurizer.blocks.20.norm1.weight": "model-00001-of-00004.safetensors",
847
+ "vision_backbone.fused_featurizer.blocks.20.norm2.bias": "model-00001-of-00004.safetensors",
848
+ "vision_backbone.fused_featurizer.blocks.20.norm2.weight": "model-00001-of-00004.safetensors",
849
+ "vision_backbone.fused_featurizer.blocks.21.attn.proj.bias": "model-00001-of-00004.safetensors",
850
+ "vision_backbone.fused_featurizer.blocks.21.attn.proj.weight": "model-00001-of-00004.safetensors",
851
+ "vision_backbone.fused_featurizer.blocks.21.attn.qkv.bias": "model-00001-of-00004.safetensors",
852
+ "vision_backbone.fused_featurizer.blocks.21.attn.qkv.weight": "model-00001-of-00004.safetensors",
853
+ "vision_backbone.fused_featurizer.blocks.21.mlp.fc1.bias": "model-00001-of-00004.safetensors",
854
+ "vision_backbone.fused_featurizer.blocks.21.mlp.fc1.weight": "model-00001-of-00004.safetensors",
855
+ "vision_backbone.fused_featurizer.blocks.21.mlp.fc2.bias": "model-00001-of-00004.safetensors",
856
+ "vision_backbone.fused_featurizer.blocks.21.mlp.fc2.weight": "model-00001-of-00004.safetensors",
857
+ "vision_backbone.fused_featurizer.blocks.21.norm1.bias": "model-00001-of-00004.safetensors",
858
+ "vision_backbone.fused_featurizer.blocks.21.norm1.weight": "model-00001-of-00004.safetensors",
859
+ "vision_backbone.fused_featurizer.blocks.21.norm2.bias": "model-00001-of-00004.safetensors",
860
+ "vision_backbone.fused_featurizer.blocks.21.norm2.weight": "model-00001-of-00004.safetensors",
861
+ "vision_backbone.fused_featurizer.blocks.22.attn.proj.bias": "model-00001-of-00004.safetensors",
862
+ "vision_backbone.fused_featurizer.blocks.22.attn.proj.weight": "model-00001-of-00004.safetensors",
863
+ "vision_backbone.fused_featurizer.blocks.22.attn.qkv.bias": "model-00001-of-00004.safetensors",
864
+ "vision_backbone.fused_featurizer.blocks.22.attn.qkv.weight": "model-00001-of-00004.safetensors",
865
+ "vision_backbone.fused_featurizer.blocks.22.mlp.fc1.bias": "model-00001-of-00004.safetensors",
866
+ "vision_backbone.fused_featurizer.blocks.22.mlp.fc1.weight": "model-00001-of-00004.safetensors",
867
+ "vision_backbone.fused_featurizer.blocks.22.mlp.fc2.bias": "model-00001-of-00004.safetensors",
868
+ "vision_backbone.fused_featurizer.blocks.22.mlp.fc2.weight": "model-00001-of-00004.safetensors",
869
+ "vision_backbone.fused_featurizer.blocks.22.norm1.bias": "model-00001-of-00004.safetensors",
870
+ "vision_backbone.fused_featurizer.blocks.22.norm1.weight": "model-00001-of-00004.safetensors",
871
+ "vision_backbone.fused_featurizer.blocks.22.norm2.bias": "model-00001-of-00004.safetensors",
872
+ "vision_backbone.fused_featurizer.blocks.22.norm2.weight": "model-00001-of-00004.safetensors",
873
+ "vision_backbone.fused_featurizer.blocks.23.attn.proj.bias": "model-00001-of-00004.safetensors",
874
+ "vision_backbone.fused_featurizer.blocks.23.attn.proj.weight": "model-00001-of-00004.safetensors",
875
+ "vision_backbone.fused_featurizer.blocks.23.attn.qkv.bias": "model-00001-of-00004.safetensors",
876
+ "vision_backbone.fused_featurizer.blocks.23.attn.qkv.weight": "model-00001-of-00004.safetensors",
877
+ "vision_backbone.fused_featurizer.blocks.23.mlp.fc1.bias": "model-00001-of-00004.safetensors",
878
+ "vision_backbone.fused_featurizer.blocks.23.mlp.fc1.weight": "model-00001-of-00004.safetensors",
879
+ "vision_backbone.fused_featurizer.blocks.23.mlp.fc2.bias": "model-00001-of-00004.safetensors",
880
+ "vision_backbone.fused_featurizer.blocks.23.mlp.fc2.weight": "model-00001-of-00004.safetensors",
881
+ "vision_backbone.fused_featurizer.blocks.23.norm1.bias": "model-00001-of-00004.safetensors",
882
+ "vision_backbone.fused_featurizer.blocks.23.norm1.weight": "model-00001-of-00004.safetensors",
883
+ "vision_backbone.fused_featurizer.blocks.23.norm2.bias": "model-00001-of-00004.safetensors",
884
+ "vision_backbone.fused_featurizer.blocks.23.norm2.weight": "model-00001-of-00004.safetensors",
885
+ "vision_backbone.fused_featurizer.blocks.24.attn.proj.bias": "model-00001-of-00004.safetensors",
886
+ "vision_backbone.fused_featurizer.blocks.24.attn.proj.weight": "model-00001-of-00004.safetensors",
887
+ "vision_backbone.fused_featurizer.blocks.24.attn.qkv.bias": "model-00001-of-00004.safetensors",
888
+ "vision_backbone.fused_featurizer.blocks.24.attn.qkv.weight": "model-00001-of-00004.safetensors",
889
+ "vision_backbone.fused_featurizer.blocks.24.mlp.fc1.bias": "model-00001-of-00004.safetensors",
890
+ "vision_backbone.fused_featurizer.blocks.24.mlp.fc1.weight": "model-00001-of-00004.safetensors",
891
+ "vision_backbone.fused_featurizer.blocks.24.mlp.fc2.bias": "model-00001-of-00004.safetensors",
892
+ "vision_backbone.fused_featurizer.blocks.24.mlp.fc2.weight": "model-00001-of-00004.safetensors",
893
+ "vision_backbone.fused_featurizer.blocks.24.norm1.bias": "model-00001-of-00004.safetensors",
894
+ "vision_backbone.fused_featurizer.blocks.24.norm1.weight": "model-00001-of-00004.safetensors",
895
+ "vision_backbone.fused_featurizer.blocks.24.norm2.bias": "model-00001-of-00004.safetensors",
896
+ "vision_backbone.fused_featurizer.blocks.24.norm2.weight": "model-00001-of-00004.safetensors",
897
+ "vision_backbone.fused_featurizer.blocks.25.attn.proj.bias": "model-00001-of-00004.safetensors",
898
+ "vision_backbone.fused_featurizer.blocks.25.attn.proj.weight": "model-00001-of-00004.safetensors",
899
+ "vision_backbone.fused_featurizer.blocks.25.attn.qkv.bias": "model-00001-of-00004.safetensors",
900
+ "vision_backbone.fused_featurizer.blocks.25.attn.qkv.weight": "model-00001-of-00004.safetensors",
901
+ "vision_backbone.fused_featurizer.blocks.25.mlp.fc1.bias": "model-00001-of-00004.safetensors",
902
+ "vision_backbone.fused_featurizer.blocks.25.mlp.fc1.weight": "model-00001-of-00004.safetensors",
903
+ "vision_backbone.fused_featurizer.blocks.25.mlp.fc2.bias": "model-00001-of-00004.safetensors",
904
+ "vision_backbone.fused_featurizer.blocks.25.mlp.fc2.weight": "model-00001-of-00004.safetensors",
905
+ "vision_backbone.fused_featurizer.blocks.25.norm1.bias": "model-00001-of-00004.safetensors",
906
+ "vision_backbone.fused_featurizer.blocks.25.norm1.weight": "model-00001-of-00004.safetensors",
907
+ "vision_backbone.fused_featurizer.blocks.25.norm2.bias": "model-00001-of-00004.safetensors",
908
+ "vision_backbone.fused_featurizer.blocks.25.norm2.weight": "model-00001-of-00004.safetensors",
909
+ "vision_backbone.fused_featurizer.blocks.26.attn.proj.bias": "model-00001-of-00004.safetensors",
910
+ "vision_backbone.fused_featurizer.blocks.26.attn.proj.weight": "model-00001-of-00004.safetensors",
911
+ "vision_backbone.fused_featurizer.blocks.26.attn.qkv.bias": "model-00001-of-00004.safetensors",
912
+ "vision_backbone.fused_featurizer.blocks.26.attn.qkv.weight": "model-00001-of-00004.safetensors",
913
+ "vision_backbone.fused_featurizer.blocks.26.mlp.fc1.bias": "model-00001-of-00004.safetensors",
914
+ "vision_backbone.fused_featurizer.blocks.26.mlp.fc1.weight": "model-00001-of-00004.safetensors",
915
+ "vision_backbone.fused_featurizer.blocks.26.mlp.fc2.bias": "model-00001-of-00004.safetensors",
916
+ "vision_backbone.fused_featurizer.blocks.26.mlp.fc2.weight": "model-00001-of-00004.safetensors",
917
+ "vision_backbone.fused_featurizer.blocks.26.norm1.bias": "model-00001-of-00004.safetensors",
918
+ "vision_backbone.fused_featurizer.blocks.26.norm1.weight": "model-00001-of-00004.safetensors",
919
+ "vision_backbone.fused_featurizer.blocks.26.norm2.bias": "model-00001-of-00004.safetensors",
920
+ "vision_backbone.fused_featurizer.blocks.26.norm2.weight": "model-00001-of-00004.safetensors",
921
+ "vision_backbone.fused_featurizer.blocks.3.attn.proj.bias": "model-00001-of-00004.safetensors",
922
+ "vision_backbone.fused_featurizer.blocks.3.attn.proj.weight": "model-00001-of-00004.safetensors",
923
+ "vision_backbone.fused_featurizer.blocks.3.attn.qkv.bias": "model-00001-of-00004.safetensors",
924
+ "vision_backbone.fused_featurizer.blocks.3.attn.qkv.weight": "model-00001-of-00004.safetensors",
925
+ "vision_backbone.fused_featurizer.blocks.3.mlp.fc1.bias": "model-00001-of-00004.safetensors",
926
+ "vision_backbone.fused_featurizer.blocks.3.mlp.fc1.weight": "model-00001-of-00004.safetensors",
927
+ "vision_backbone.fused_featurizer.blocks.3.mlp.fc2.bias": "model-00001-of-00004.safetensors",
928
+ "vision_backbone.fused_featurizer.blocks.3.mlp.fc2.weight": "model-00001-of-00004.safetensors",
929
+ "vision_backbone.fused_featurizer.blocks.3.norm1.bias": "model-00001-of-00004.safetensors",
930
+ "vision_backbone.fused_featurizer.blocks.3.norm1.weight": "model-00001-of-00004.safetensors",
931
+ "vision_backbone.fused_featurizer.blocks.3.norm2.bias": "model-00001-of-00004.safetensors",
932
+ "vision_backbone.fused_featurizer.blocks.3.norm2.weight": "model-00001-of-00004.safetensors",
933
+ "vision_backbone.fused_featurizer.blocks.4.attn.proj.bias": "model-00001-of-00004.safetensors",
934
+ "vision_backbone.fused_featurizer.blocks.4.attn.proj.weight": "model-00001-of-00004.safetensors",
935
+ "vision_backbone.fused_featurizer.blocks.4.attn.qkv.bias": "model-00001-of-00004.safetensors",
936
+ "vision_backbone.fused_featurizer.blocks.4.attn.qkv.weight": "model-00001-of-00004.safetensors",
937
+ "vision_backbone.fused_featurizer.blocks.4.mlp.fc1.bias": "model-00001-of-00004.safetensors",
938
+ "vision_backbone.fused_featurizer.blocks.4.mlp.fc1.weight": "model-00001-of-00004.safetensors",
939
+ "vision_backbone.fused_featurizer.blocks.4.mlp.fc2.bias": "model-00001-of-00004.safetensors",
940
+ "vision_backbone.fused_featurizer.blocks.4.mlp.fc2.weight": "model-00001-of-00004.safetensors",
941
+ "vision_backbone.fused_featurizer.blocks.4.norm1.bias": "model-00001-of-00004.safetensors",
942
+ "vision_backbone.fused_featurizer.blocks.4.norm1.weight": "model-00001-of-00004.safetensors",
943
+ "vision_backbone.fused_featurizer.blocks.4.norm2.bias": "model-00001-of-00004.safetensors",
944
+ "vision_backbone.fused_featurizer.blocks.4.norm2.weight": "model-00001-of-00004.safetensors",
945
+ "vision_backbone.fused_featurizer.blocks.5.attn.proj.bias": "model-00001-of-00004.safetensors",
946
+ "vision_backbone.fused_featurizer.blocks.5.attn.proj.weight": "model-00001-of-00004.safetensors",
947
+ "vision_backbone.fused_featurizer.blocks.5.attn.qkv.bias": "model-00001-of-00004.safetensors",
948
+ "vision_backbone.fused_featurizer.blocks.5.attn.qkv.weight": "model-00001-of-00004.safetensors",
949
+ "vision_backbone.fused_featurizer.blocks.5.mlp.fc1.bias": "model-00001-of-00004.safetensors",
950
+ "vision_backbone.fused_featurizer.blocks.5.mlp.fc1.weight": "model-00001-of-00004.safetensors",
951
+ "vision_backbone.fused_featurizer.blocks.5.mlp.fc2.bias": "model-00001-of-00004.safetensors",
952
+ "vision_backbone.fused_featurizer.blocks.5.mlp.fc2.weight": "model-00001-of-00004.safetensors",
953
+ "vision_backbone.fused_featurizer.blocks.5.norm1.bias": "model-00001-of-00004.safetensors",
954
+ "vision_backbone.fused_featurizer.blocks.5.norm1.weight": "model-00001-of-00004.safetensors",
955
+ "vision_backbone.fused_featurizer.blocks.5.norm2.bias": "model-00001-of-00004.safetensors",
956
+ "vision_backbone.fused_featurizer.blocks.5.norm2.weight": "model-00001-of-00004.safetensors",
957
+ "vision_backbone.fused_featurizer.blocks.6.attn.proj.bias": "model-00001-of-00004.safetensors",
958
+ "vision_backbone.fused_featurizer.blocks.6.attn.proj.weight": "model-00001-of-00004.safetensors",
959
+ "vision_backbone.fused_featurizer.blocks.6.attn.qkv.bias": "model-00001-of-00004.safetensors",
960
+ "vision_backbone.fused_featurizer.blocks.6.attn.qkv.weight": "model-00001-of-00004.safetensors",
961
+ "vision_backbone.fused_featurizer.blocks.6.mlp.fc1.bias": "model-00001-of-00004.safetensors",
962
+ "vision_backbone.fused_featurizer.blocks.6.mlp.fc1.weight": "model-00001-of-00004.safetensors",
963
+ "vision_backbone.fused_featurizer.blocks.6.mlp.fc2.bias": "model-00001-of-00004.safetensors",
964
+ "vision_backbone.fused_featurizer.blocks.6.mlp.fc2.weight": "model-00001-of-00004.safetensors",
965
+ "vision_backbone.fused_featurizer.blocks.6.norm1.bias": "model-00001-of-00004.safetensors",
966
+ "vision_backbone.fused_featurizer.blocks.6.norm1.weight": "model-00001-of-00004.safetensors",
967
+ "vision_backbone.fused_featurizer.blocks.6.norm2.bias": "model-00001-of-00004.safetensors",
968
+ "vision_backbone.fused_featurizer.blocks.6.norm2.weight": "model-00001-of-00004.safetensors",
969
+ "vision_backbone.fused_featurizer.blocks.7.attn.proj.bias": "model-00001-of-00004.safetensors",
970
+ "vision_backbone.fused_featurizer.blocks.7.attn.proj.weight": "model-00001-of-00004.safetensors",
971
+ "vision_backbone.fused_featurizer.blocks.7.attn.qkv.bias": "model-00001-of-00004.safetensors",
972
+ "vision_backbone.fused_featurizer.blocks.7.attn.qkv.weight": "model-00001-of-00004.safetensors",
973
+ "vision_backbone.fused_featurizer.blocks.7.mlp.fc1.bias": "model-00001-of-00004.safetensors",
974
+ "vision_backbone.fused_featurizer.blocks.7.mlp.fc1.weight": "model-00001-of-00004.safetensors",
975
+ "vision_backbone.fused_featurizer.blocks.7.mlp.fc2.bias": "model-00001-of-00004.safetensors",
976
+ "vision_backbone.fused_featurizer.blocks.7.mlp.fc2.weight": "model-00001-of-00004.safetensors",
977
+ "vision_backbone.fused_featurizer.blocks.7.norm1.bias": "model-00001-of-00004.safetensors",
978
+ "vision_backbone.fused_featurizer.blocks.7.norm1.weight": "model-00001-of-00004.safetensors",
979
+ "vision_backbone.fused_featurizer.blocks.7.norm2.bias": "model-00001-of-00004.safetensors",
980
+ "vision_backbone.fused_featurizer.blocks.7.norm2.weight": "model-00001-of-00004.safetensors",
981
+ "vision_backbone.fused_featurizer.blocks.8.attn.proj.bias": "model-00001-of-00004.safetensors",
982
+ "vision_backbone.fused_featurizer.blocks.8.attn.proj.weight": "model-00001-of-00004.safetensors",
983
+ "vision_backbone.fused_featurizer.blocks.8.attn.qkv.bias": "model-00001-of-00004.safetensors",
984
+ "vision_backbone.fused_featurizer.blocks.8.attn.qkv.weight": "model-00001-of-00004.safetensors",
985
+ "vision_backbone.fused_featurizer.blocks.8.mlp.fc1.bias": "model-00001-of-00004.safetensors",
986
+ "vision_backbone.fused_featurizer.blocks.8.mlp.fc1.weight": "model-00001-of-00004.safetensors",
987
+ "vision_backbone.fused_featurizer.blocks.8.mlp.fc2.bias": "model-00001-of-00004.safetensors",
988
+ "vision_backbone.fused_featurizer.blocks.8.mlp.fc2.weight": "model-00001-of-00004.safetensors",
989
+ "vision_backbone.fused_featurizer.blocks.8.norm1.bias": "model-00001-of-00004.safetensors",
990
+ "vision_backbone.fused_featurizer.blocks.8.norm1.weight": "model-00001-of-00004.safetensors",
991
+ "vision_backbone.fused_featurizer.blocks.8.norm2.bias": "model-00001-of-00004.safetensors",
992
+ "vision_backbone.fused_featurizer.blocks.8.norm2.weight": "model-00001-of-00004.safetensors",
993
+ "vision_backbone.fused_featurizer.blocks.9.attn.proj.bias": "model-00001-of-00004.safetensors",
994
+ "vision_backbone.fused_featurizer.blocks.9.attn.proj.weight": "model-00001-of-00004.safetensors",
995
+ "vision_backbone.fused_featurizer.blocks.9.attn.qkv.bias": "model-00001-of-00004.safetensors",
996
+ "vision_backbone.fused_featurizer.blocks.9.attn.qkv.weight": "model-00001-of-00004.safetensors",
997
+ "vision_backbone.fused_featurizer.blocks.9.mlp.fc1.bias": "model-00001-of-00004.safetensors",
998
+ "vision_backbone.fused_featurizer.blocks.9.mlp.fc1.weight": "model-00001-of-00004.safetensors",
999
+ "vision_backbone.fused_featurizer.blocks.9.mlp.fc2.bias": "model-00001-of-00004.safetensors",
1000
+ "vision_backbone.fused_featurizer.blocks.9.mlp.fc2.weight": "model-00001-of-00004.safetensors",
1001
+ "vision_backbone.fused_featurizer.blocks.9.norm1.bias": "model-00001-of-00004.safetensors",
1002
+ "vision_backbone.fused_featurizer.blocks.9.norm1.weight": "model-00001-of-00004.safetensors",
1003
+ "vision_backbone.fused_featurizer.blocks.9.norm2.bias": "model-00001-of-00004.safetensors",
1004
+ "vision_backbone.fused_featurizer.blocks.9.norm2.weight": "model-00001-of-00004.safetensors",
1005
+ "vision_backbone.fused_featurizer.norm.bias": "model-00001-of-00004.safetensors",
1006
+ "vision_backbone.fused_featurizer.norm.weight": "model-00001-of-00004.safetensors",
1007
+ "vision_backbone.fused_featurizer.patch_embed.proj.bias": "model-00001-of-00004.safetensors",
1008
+ "vision_backbone.fused_featurizer.patch_embed.proj.weight": "model-00001-of-00004.safetensors",
1009
+ "vision_backbone.fused_featurizer.pos_embed": "model-00001-of-00004.safetensors"
1010
+ }
1011
+ }
modeling_prismatic.py ADDED
@@ -0,0 +1,1788 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ modeling_prismatic.py
3
+
4
+ Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.
5
+ Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,
6
+ but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.
7
+ """
8
+
9
+ import logging
10
+ import math
11
+ from dataclasses import dataclass
12
+ from functools import partial
13
+ from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
14
+
15
+ import numpy as np
16
+ import timm
17
+ import tokenizers
18
+ import torch
19
+ import torch.nn as nn
20
+ import transformers
21
+ from timm.models.vision_transformer import LayerScale
22
+ from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
23
+ from transformers.modeling_outputs import ModelOutput
24
+
25
+ from prismatic.training.train_utils import (
26
+ get_current_action_mask,
27
+ get_next_actions_mask,
28
+ )
29
+ from prismatic.vla.constants import (
30
+ ACTION_DIM,
31
+ ACTION_PROPRIO_NORMALIZATION_TYPE,
32
+ ACTION_TOKEN_BEGIN_IDX,
33
+ IGNORE_INDEX,
34
+ NUM_ACTIONS_CHUNK,
35
+ STOP_INDEX,
36
+ NormalizationType,
37
+ )
38
+
39
+ from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
40
+
41
+ # Set up logger
42
+ logger = logging.getLogger(__name__)
43
+
44
+
45
+ # === Utility Functions for Monkey-Patching ===
46
+ def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
47
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
48
+ result = fn(*args, **kwargs)
49
+ return result[0] if isinstance(result, tuple) else result
50
+
51
+ return wrapper
52
+
53
+
54
+ # HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
55
+ # =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
56
+ # =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
57
+ def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
58
+ return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
59
+
60
+
61
+ def ls_apply_patch(ls_module: LayerScale):
62
+ ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
63
+ ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
64
+ del ls_module.gamma
65
+
66
+
67
+ # === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
68
+ class PrismaticVisionBackbone(nn.Module):
69
+ """
70
+ Vision backbone for Prismatic models that handles image feature extraction.
71
+
72
+ Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
73
+ For fused backbones, features from both models are concatenated along the feature dimension.
74
+ """
75
+
76
+ def __init__(
77
+ self,
78
+ use_fused_vision_backbone: bool,
79
+ image_sizes: List[int],
80
+ timm_model_ids: List[str],
81
+ timm_override_act_layers: List[Optional[str]],
82
+ ) -> None:
83
+ """
84
+ Initialize the vision backbone.
85
+
86
+ Args:
87
+ use_fused_vision_backbone: Whether to use two backbones and fuse their features
88
+ image_sizes: List of image sizes for each backbone
89
+ timm_model_ids: List of TIMM model IDs to use for each backbone
90
+ timm_override_act_layers: List of activation layer overrides for each backbone
91
+ """
92
+ super().__init__()
93
+ self.use_fused_vision_backbone = use_fused_vision_backbone
94
+ self.num_images_in_input = 1 # Default value, can be overridden later
95
+
96
+ # Validate number of (fused) vision backbones
97
+ if len(timm_model_ids) > 2:
98
+ raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
99
+
100
+ # Create primary featurizer
101
+ self.featurizer = self._create_featurizer(
102
+ model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
103
+ )
104
+ self.embed_dim = self.featurizer.embed_dim
105
+
106
+ # Create secondary featurizer if using fused backbone
107
+ if self.use_fused_vision_backbone:
108
+ self.fused_featurizer = self._create_featurizer(
109
+ model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
110
+ )
111
+ self.embed_dim += self.fused_featurizer.embed_dim
112
+
113
+ # Patch LayerScale modules for HF compatibility
114
+ self._patch_layer_scales()
115
+
116
+ def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
117
+ """
118
+ Create a TIMM-based featurizer model with appropriate configurations.
119
+
120
+ Args:
121
+ model_id: The TIMM model ID to load
122
+ img_size: Input image size for the model
123
+ act_layer: Override for the activation layer type
124
+
125
+ Returns:
126
+ A configured featurizer model
127
+ """
128
+ featurizer = timm.create_model(
129
+ model_id,
130
+ pretrained=False,
131
+ num_classes=0,
132
+ img_size=img_size,
133
+ act_layer=act_layer,
134
+ )
135
+
136
+ # Monkey-patch the forward function to extract the second-to-last layer features
137
+ num_blocks = len(featurizer.blocks)
138
+ featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
139
+
140
+ return featurizer
141
+
142
+ def _patch_layer_scales(self) -> None:
143
+ """
144
+ Patch all LayerScale modules to be compatible with HF's parameter naming.
145
+
146
+ HF Transformers overwrites parameters with names containing 'gamma',
147
+ so we need to rename and modify the forward method.
148
+ """
149
+ # Patch primary featurizer
150
+ for module in self.featurizer.modules():
151
+ if isinstance(module, LayerScale):
152
+ ls_apply_patch(module)
153
+
154
+ # Patch secondary featurizer if it exists
155
+ if self.use_fused_vision_backbone:
156
+ for module in self.fused_featurizer.modules():
157
+ if isinstance(module, LayerScale):
158
+ ls_apply_patch(module)
159
+
160
+ def get_num_patches(self) -> int:
161
+ """
162
+ Returns the number of vision patches output by the vision backbone.
163
+
164
+ Returns:
165
+ Number of patches per image
166
+ """
167
+ return self.featurizer.patch_embed.num_patches
168
+
169
+ def get_num_images_in_input(self) -> int:
170
+ """
171
+ Returns the number of input images for the vision backbone.
172
+
173
+ Returns:
174
+ Number of images expected in the input
175
+ """
176
+ return self.num_images_in_input
177
+
178
+ def set_num_images_in_input(self, num_images_in_input: int) -> None:
179
+ """
180
+ Sets the number of input images for the vision backbone.
181
+
182
+ Args:
183
+ num_images_in_input: Number of images to expect in the input
184
+ """
185
+ self.num_images_in_input = num_images_in_input
186
+
187
+ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
188
+ """
189
+ Implements the forward pass for the vision backbone.
190
+
191
+ If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
192
+ (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
193
+
194
+ Args:
195
+ pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
196
+ """
197
+ if self.num_images_in_input == 1:
198
+ if not self.use_fused_vision_backbone:
199
+ return self.featurizer(pixel_values)
200
+
201
+ # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
202
+ img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
203
+ patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
204
+
205
+ return torch.cat([patches, patches_fused], dim=2)
206
+
207
+ else:
208
+ assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
209
+
210
+ # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
211
+ images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
212
+
213
+ # Process each image and collect patches
214
+ all_patches = []
215
+ for img in images:
216
+ # Split each image further into two stacks of channels (each with 3 channels)
217
+ img_regular, img_fused = torch.split(img, [3, 3], dim=1)
218
+
219
+ # Get patches from both SigLIP and DINOv2 vision transformers
220
+ patches = self.featurizer(img_regular)
221
+ patches_fused = self.fused_featurizer(img_fused)
222
+
223
+ # Concatenate SigLIP and DINOv2 patches along the hidden dimension
224
+ combined_patches = torch.cat([patches, patches_fused], dim=2)
225
+ all_patches.append(combined_patches)
226
+
227
+ # Concatenate all patches along the patch dimension
228
+ return torch.cat(all_patches, dim=1)
229
+
230
+
231
+ # === Prismatic Projector (nn.Module) Definitions ===
232
+ class PrismaticProjector(nn.Module):
233
+ def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
234
+ super().__init__()
235
+ self.use_fused_vision_backbone = use_fused_vision_backbone
236
+ self.vision_dim, self.llm_dim = vision_dim, llm_dim
237
+
238
+ # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
239
+ if not self.use_fused_vision_backbone:
240
+ self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
241
+ self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
242
+ self.act_fn1 = nn.GELU()
243
+ else:
244
+ initial_projection_dim = 4 * vision_dim
245
+ self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
246
+ self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
247
+ self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
248
+ self.act_fn1 = nn.GELU()
249
+ self.act_fn2 = nn.GELU()
250
+
251
+ def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
252
+ if not self.use_fused_vision_backbone:
253
+ projected_features = self.fc1(img_patches)
254
+ projected_features = self.act_fn1(projected_features)
255
+ projected_features = self.fc2(projected_features)
256
+ else:
257
+ projected_features = self.fc1(img_patches)
258
+ projected_features = self.act_fn1(projected_features)
259
+ projected_features = self.fc2(projected_features)
260
+ projected_features = self.act_fn2(projected_features)
261
+ projected_features = self.fc3(projected_features)
262
+
263
+ return projected_features
264
+
265
+
266
+ class GazingSelfAttentionBlock(nn.Module):
267
+ def __init__(self, dim: int, num_heads: int, ffn_mult: int = 4):
268
+ super().__init__()
269
+ self.attn_norm = nn.LayerNorm(dim)
270
+ self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True, dropout=0.0)
271
+ self.ffn_norm = nn.LayerNorm(dim)
272
+ self.ffn = nn.Sequential(
273
+ nn.Linear(dim, ffn_mult * dim),
274
+ nn.GELU(),
275
+ nn.Linear(ffn_mult * dim, dim),
276
+ )
277
+
278
+ def forward(self, tokens: torch.Tensor) -> torch.Tensor:
279
+ attn_input = self.attn_norm(tokens)
280
+ attn_output, _ = self.attn(attn_input, attn_input, attn_input, need_weights=False)
281
+ tokens = tokens + attn_output
282
+ tokens = tokens + self.ffn(self.ffn_norm(tokens))
283
+ return tokens
284
+
285
+
286
+ class TextConditionedTokenGate(nn.Module):
287
+ def __init__(
288
+ self,
289
+ llm_dim: int,
290
+ use_text_summary: bool = True,
291
+ use_vision_tokens: bool = True,
292
+ hidden_dim: int = 512,
293
+ text_pooling_mode: str = "mean",
294
+ text_pool_hidden_dim: int = 128,
295
+ cross_attention_dim: int = 256,
296
+ cross_attention_heads: int = 1,
297
+ gate_mlp_depth: int = 1,
298
+ contrastive_visual_tau: float = 0.1,
299
+ contrastive_text_tau: float = 1.0,
300
+ gazing_mode: str = "mlp",
301
+ self_attn_dim: int = 512,
302
+ self_attn_heads: int = 8,
303
+ self_attn_layers: int = 1,
304
+ ):
305
+ super().__init__()
306
+ if text_pooling_mode not in {"mean", "mlp", "cross_attention", "contrastive_alignment_score"}:
307
+ raise ValueError(
308
+ "`text_pooling_mode` must be one of "
309
+ "['mean', 'mlp', 'cross_attention', 'contrastive_alignment_score'], "
310
+ f"got {text_pooling_mode!r}"
311
+ )
312
+ if gate_mlp_depth <= 0:
313
+ raise ValueError(f"`gate_mlp_depth` must be positive, got {gate_mlp_depth}")
314
+ if contrastive_visual_tau <= 0:
315
+ raise ValueError(f"`contrastive_visual_tau` must be positive, got {contrastive_visual_tau}")
316
+ if contrastive_text_tau <= 0:
317
+ raise ValueError(f"`contrastive_text_tau` must be positive, got {contrastive_text_tau}")
318
+ if gazing_mode not in {"mlp", "self_attention"}:
319
+ raise ValueError(f"`gazing_mode` must be one of ['mlp', 'self_attention'], got {gazing_mode!r}")
320
+ if self_attn_dim <= 0:
321
+ raise ValueError(f"`self_attn_dim` must be positive, got {self_attn_dim}")
322
+ if self_attn_heads <= 0:
323
+ raise ValueError(f"`self_attn_heads` must be positive, got {self_attn_heads}")
324
+ if self_attn_dim % self_attn_heads != 0:
325
+ raise ValueError(
326
+ "`self_attn_dim` must be divisible by `self_attn_heads`; "
327
+ f"got {self_attn_dim} and {self_attn_heads}"
328
+ )
329
+ if self_attn_layers <= 0:
330
+ raise ValueError(f"`self_attn_layers` must be positive, got {self_attn_layers}")
331
+ self.llm_dim = llm_dim
332
+ self.use_text_summary = use_text_summary
333
+ self.use_vision_tokens = use_vision_tokens
334
+ self.hidden_dim = hidden_dim
335
+ self.text_pooling_mode = text_pooling_mode
336
+ self.text_pool_hidden_dim = text_pool_hidden_dim
337
+ self.cross_attention_dim = cross_attention_dim
338
+ self.cross_attention_heads = cross_attention_heads
339
+ self.cross_attention_head_dim = None
340
+ self.gate_mlp_depth = gate_mlp_depth
341
+ self.contrastive_visual_tau = contrastive_visual_tau
342
+ self.contrastive_text_tau = contrastive_text_tau
343
+ self.gazing_mode = gazing_mode
344
+ self.self_attn_dim = self_attn_dim
345
+ self.self_attn_heads = self_attn_heads
346
+ self.self_attn_layers = self_attn_layers
347
+ self.text_pool = None
348
+ self.text_query = None
349
+ self.text_key = None
350
+ self.text_value = None
351
+ self.text_output = None
352
+ self.visual_attn_ln = None
353
+ self.text_attn_ln = None
354
+ self.gate = None
355
+ self.context_proj = None
356
+ self.gazing_blocks = nn.ModuleList()
357
+ self.gate_out = None
358
+ self.gazing_output_ln = None
359
+ self.gate_visual_ln = nn.LayerNorm(llm_dim)
360
+ self.gate_text_ln = nn.LayerNorm(llm_dim)
361
+ self.last_text_pool_mask = None
362
+ self.last_contrastive_scores = None
363
+ gate_input_dim = 2 * llm_dim if self.use_text_summary and self.use_vision_tokens else llm_dim
364
+ if self.text_pooling_mode == "mlp":
365
+ self.text_pool = nn.Sequential(
366
+ nn.Linear(llm_dim, text_pool_hidden_dim),
367
+ nn.Tanh(),
368
+ nn.Linear(text_pool_hidden_dim, 1),
369
+ )
370
+ elif self.text_pooling_mode == "cross_attention":
371
+ if self.cross_attention_heads <= 0:
372
+ raise ValueError(
373
+ "`cross_attention_heads` must be positive, "
374
+ f"got {self.cross_attention_heads}"
375
+ )
376
+ if self.cross_attention_dim % self.cross_attention_heads != 0:
377
+ raise ValueError(
378
+ "`cross_attention_dim` must be divisible by `cross_attention_heads`; got "
379
+ f"{self.cross_attention_dim} and {self.cross_attention_heads}"
380
+ )
381
+ self.cross_attention_head_dim = self.cross_attention_dim // self.cross_attention_heads
382
+ self.visual_attn_ln = nn.LayerNorm(llm_dim)
383
+ self.text_attn_ln = nn.LayerNorm(llm_dim)
384
+ self.text_query = nn.Linear(llm_dim, cross_attention_dim, bias=False)
385
+ self.text_key = nn.Linear(llm_dim, cross_attention_dim, bias=False)
386
+ self.text_value = nn.Linear(llm_dim, cross_attention_dim, bias=False)
387
+ self.text_output = nn.Linear(cross_attention_dim, llm_dim, bias=False)
388
+
389
+ if self.gazing_mode == "mlp":
390
+ gate_layers = [nn.Linear(gate_input_dim, hidden_dim), nn.GELU()]
391
+ for _ in range(gate_mlp_depth - 1):
392
+ gate_layers.extend([nn.Linear(hidden_dim, hidden_dim), nn.GELU()])
393
+ gate_layers.append(nn.Linear(hidden_dim, 1))
394
+ self.gate = nn.Sequential(*gate_layers)
395
+ else:
396
+ self.context_proj = nn.Linear(gate_input_dim, self_attn_dim)
397
+ self.gazing_blocks = nn.ModuleList(
398
+ [GazingSelfAttentionBlock(self_attn_dim, self_attn_heads) for _ in range(self_attn_layers)]
399
+ )
400
+ self.gazing_output_ln = nn.LayerNorm(self_attn_dim)
401
+ self.gate_out = nn.Linear(self_attn_dim, 1)
402
+
403
+ @staticmethod
404
+ def _pool_text_embeddings(
405
+ language_embeddings: torch.Tensor, language_attention_mask: Optional[torch.Tensor] = None
406
+ ) -> torch.Tensor:
407
+ if language_attention_mask is None:
408
+ return language_embeddings.mean(dim=1)
409
+
410
+ text_mask = language_attention_mask.to(dtype=language_embeddings.dtype).unsqueeze(-1) # [B, T, 1]
411
+ valid_token_counts = text_mask.sum(dim=1).clamp(min=1.0) # [B, 1]
412
+ masked_language_embeddings = language_embeddings * text_mask
413
+ return masked_language_embeddings.sum(dim=1) / valid_token_counts
414
+
415
+ def forward(
416
+ self,
417
+ visual_tokens: torch.Tensor,
418
+ language_embeddings: torch.Tensor,
419
+ language_attention_mask: Optional[torch.Tensor] = None,
420
+ language_text_pool_mask: Optional[torch.Tensor] = None,
421
+ patches_per_image: Optional[int] = None,
422
+ ):
423
+ # visual_tokens: [B, N, D]
424
+ # language_embeddings: [B, T, D]
425
+ self.last_text_pool_mask = None
426
+ self.last_contrastive_scores = None
427
+ text_pool_mask = language_text_pool_mask
428
+ if text_pool_mask is not None:
429
+ text_pool_mask = text_pool_mask.to(device=language_embeddings.device, dtype=torch.bool)
430
+ if language_attention_mask is not None:
431
+ text_pool_mask = text_pool_mask & language_attention_mask.to(device=language_embeddings.device).bool()
432
+ has_pool_token = text_pool_mask.any(dim=1, keepdim=True)
433
+ if not has_pool_token.all():
434
+ fallback_mask = (
435
+ language_attention_mask.to(device=language_embeddings.device).bool()
436
+ if language_attention_mask is not None
437
+ else torch.ones_like(text_pool_mask, dtype=torch.bool)
438
+ )
439
+ logger.warning(
440
+ "text pool mask has no valid tokens for at least one sample; falling back to attention mask."
441
+ )
442
+ text_pool_mask = torch.where(has_pool_token, text_pool_mask, fallback_mask)
443
+ else:
444
+ text_pool_mask = language_attention_mask
445
+ if text_pool_mask is not None:
446
+ text_pool_mask = text_pool_mask.to(device=language_embeddings.device, dtype=torch.bool)
447
+
448
+ if text_pool_mask is not None:
449
+ self.last_text_pool_mask = text_pool_mask.detach().bool()
450
+
451
+ text_pool_weights = None
452
+ if not self.use_text_summary:
453
+ gate_input = self.gate_visual_ln(visual_tokens)
454
+ elif self.text_pooling_mode == "mean":
455
+ # Pool only valid language tokens so padding does not leak into the text summary.
456
+ pooled_text = self._pool_text_embeddings(language_embeddings, text_pool_mask) # [B, D]
457
+ pooled_text = pooled_text.unsqueeze(1).expand(-1, visual_tokens.shape[1], -1) # [B, N, D]
458
+ elif self.text_pooling_mode == "mlp":
459
+ text_pool_scores = self.text_pool(language_embeddings).squeeze(-1) # [B, T]
460
+ if text_pool_mask is not None:
461
+ text_pool_scores = text_pool_scores.masked_fill(~text_pool_mask.bool(), -1e9)
462
+ text_pool_weights = torch.softmax(text_pool_scores, dim=1) # [B, T]
463
+ if text_pool_mask is not None:
464
+ text_pool_weights = text_pool_weights * text_pool_mask.to(dtype=text_pool_weights.dtype)
465
+ text_pool_weights = text_pool_weights / text_pool_weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
466
+ pooled_text = torch.sum(text_pool_weights.unsqueeze(-1) * language_embeddings, dim=1) # [B, D]
467
+ pooled_text = pooled_text.unsqueeze(1).expand(-1, visual_tokens.shape[1], -1) # [B, N, D]
468
+ elif self.text_pooling_mode == "cross_attention":
469
+ attn_visual_tokens = self.visual_attn_ln(visual_tokens)
470
+ attn_language_embeddings = self.text_attn_ln(language_embeddings)
471
+ text_queries = self.text_query(attn_visual_tokens) # [B, N, A]
472
+ text_keys = self.text_key(attn_language_embeddings) # [B, T, A]
473
+ text_values = self.text_value(attn_language_embeddings) # [B, T, A]
474
+ batch_size, num_visual_tokens, _ = text_queries.shape
475
+ num_text_tokens = text_keys.shape[1]
476
+ num_heads, head_dim = self.cross_attention_heads, self.cross_attention_head_dim
477
+ text_queries = text_queries.view(batch_size, num_visual_tokens, num_heads, head_dim).transpose(1, 2)
478
+ text_keys = text_keys.view(batch_size, num_text_tokens, num_heads, head_dim).transpose(1, 2)
479
+ text_values = text_values.view(batch_size, num_text_tokens, num_heads, head_dim).transpose(1, 2)
480
+ text_pool_scores = torch.matmul(text_queries, text_keys.transpose(-1, -2)) / math.sqrt(
481
+ head_dim
482
+ ) # [B, H, N, T]
483
+ if text_pool_mask is not None:
484
+ expanded_text_pool_mask = text_pool_mask.unsqueeze(1).unsqueeze(1).bool()
485
+ text_pool_scores = text_pool_scores.masked_fill(~expanded_text_pool_mask, -1e9)
486
+ text_pool_weights = torch.softmax(text_pool_scores, dim=-1) # [B, H, N, T]
487
+ if text_pool_mask is not None:
488
+ text_pool_weights = text_pool_weights * text_pool_mask.unsqueeze(1).unsqueeze(1).to(
489
+ dtype=text_pool_weights.dtype
490
+ )
491
+ text_pool_weights = text_pool_weights / text_pool_weights.sum(dim=-1, keepdim=True).clamp(min=1e-6)
492
+ pooled_text = torch.matmul(text_pool_weights, text_values) # [B, H, N, Dh]
493
+ pooled_text = pooled_text.transpose(1, 2).contiguous().view(
494
+ batch_size, num_visual_tokens, self.cross_attention_dim
495
+ )
496
+ pooled_text = self.text_output(pooled_text) # [B, N, D]
497
+ text_pool_weights = text_pool_weights.mean(dim=1) # [B, N, T]
498
+ else:
499
+ visual_norm = torch.nn.functional.normalize(visual_tokens.float(), dim=-1)
500
+ text_norm = torch.nn.functional.normalize(language_embeddings.float(), dim=-1)
501
+ similarity = torch.matmul(text_norm, visual_norm.transpose(1, 2)) # [B, T, N]
502
+
503
+ visual_weights = torch.softmax(similarity / self.contrastive_visual_tau, dim=-1)
504
+ local_evidence = torch.sum(visual_weights * similarity, dim=-1) # [B, T]
505
+ global_baseline = similarity.mean(dim=-1) # [B, T]
506
+ contrastive_scores = local_evidence - global_baseline
507
+ self.last_contrastive_scores = contrastive_scores.detach()
508
+
509
+ if text_pool_mask is not None:
510
+ contrastive_scores = contrastive_scores.masked_fill(~text_pool_mask.bool(), -1e9)
511
+ text_pool_weights = torch.softmax(contrastive_scores / self.contrastive_text_tau, dim=1) # [B, T]
512
+ if text_pool_mask is not None:
513
+ text_pool_weights = text_pool_weights * text_pool_mask.to(dtype=text_pool_weights.dtype)
514
+ text_pool_weights = text_pool_weights / text_pool_weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
515
+
516
+ pooled_text = torch.sum(
517
+ text_pool_weights.to(dtype=language_embeddings.dtype).unsqueeze(-1) * language_embeddings,
518
+ dim=1,
519
+ ) # [B, D]
520
+ pooled_text = pooled_text.unsqueeze(1).expand(-1, visual_tokens.shape[1], -1) # [B, N, D]
521
+
522
+ if self.use_text_summary:
523
+ gate_text_tokens = self.gate_text_ln(pooled_text)
524
+ if self.use_vision_tokens:
525
+ gate_visual_tokens = self.gate_visual_ln(visual_tokens)
526
+ gate_input = torch.cat([gate_visual_tokens, gate_text_tokens], dim=-1) # [B, N, 2D]
527
+ else:
528
+ gate_input = gate_text_tokens # [B, N, D]
529
+ if self.gazing_mode == "mlp":
530
+ gate_logits = self.gate(gate_input) # [B, N, 1]
531
+ else:
532
+ gazing_tokens = self.context_proj(gate_input)
533
+ gazing_tokens = self._apply_gazing_blocks(gazing_tokens, patches_per_image)
534
+ gazing_tokens = self.gazing_output_ln(gazing_tokens)
535
+ gate_logits = self.gate_out(gazing_tokens) # [B, N, 1]
536
+ token_gate = torch.sigmoid(gate_logits) # [B, N, 1]
537
+ gated_visual_tokens = visual_tokens * token_gate
538
+ return gated_visual_tokens, token_gate, text_pool_weights
539
+
540
+ def _apply_gazing_blocks(self, tokens: torch.Tensor, patches_per_image: Optional[int]) -> torch.Tensor:
541
+ if patches_per_image is None or patches_per_image <= 0 or tokens.shape[1] % patches_per_image != 0:
542
+ if patches_per_image is not None:
543
+ logger.warning(
544
+ "Invalid patches_per_image for self-attention gazing; falling back to full visual sequence. "
545
+ f"tokens={tokens.shape[1]}, patches_per_image={patches_per_image}"
546
+ )
547
+ for block in self.gazing_blocks:
548
+ tokens = block(tokens)
549
+ return tokens
550
+
551
+ batch_size, num_visual_tokens, dim = tokens.shape
552
+ num_images = num_visual_tokens // patches_per_image
553
+ tokens = tokens.reshape(batch_size * num_images, patches_per_image, dim)
554
+ for block in self.gazing_blocks:
555
+ tokens = block(tokens)
556
+ return tokens.reshape(batch_size, num_images * patches_per_image, dim)
557
+
558
+
559
+ # === Main HF Class Definitions ===
560
+ @dataclass
561
+ class PrismaticCausalLMOutputWithPast(ModelOutput):
562
+ """Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
563
+
564
+ loss: Optional[torch.FloatTensor] = None
565
+ logits: torch.FloatTensor = None
566
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
567
+ hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
568
+ attentions: Optional[Tuple[torch.FloatTensor]] = None
569
+
570
+ # Additions for VLMs
571
+ projector_features: Optional[torch.FloatTensor] = None
572
+
573
+
574
+ class PrismaticPreTrainedModel(PreTrainedModel):
575
+ config_class: PretrainedConfig = PrismaticConfig
576
+ base_model_prefix: str = "model"
577
+ supports_gradient_checkpointing: bool = True
578
+
579
+ _no_split_modules: ClassVar[List[str]] = ["PrismaticProjector", "TextConditionedTokenGate"]
580
+ _skip_keys_device_placement: str = "past_key_values"
581
+ _supports_flash_attn_2: bool = True
582
+
583
+ def _init_weights(self, module: nn.Module) -> None:
584
+ # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
585
+ # => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
586
+ # https://github.com/TRI-ML/prismatic-vlms
587
+ std = (
588
+ self.config.initializer_range
589
+ if hasattr(self.config, "initializer_range")
590
+ else self.config.text_config.initializer_range
591
+ )
592
+
593
+ if hasattr(module, "class_embedding"):
594
+ module.class_embedding.data.normal_(mean=0.0, std=std)
595
+
596
+ if isinstance(module, (nn.Linear, nn.Conv2d)):
597
+ module.weight.data.normal_(mean=0.0, std=std)
598
+ if module.bias is not None:
599
+ module.bias.data.zero_()
600
+ elif isinstance(module, nn.Embedding):
601
+ module.weight.data.normal_(mean=0.0, std=std)
602
+ if module.padding_idx is not None:
603
+ module.weight.data[module.padding_idx].zero_()
604
+
605
+ @property
606
+ def _supports_sdpa(self) -> bool:
607
+ """Check LLM supports SDPA Attention"""
608
+ return self.language_model._supports_sdpa
609
+
610
+
611
+ class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
612
+ def __init__(self, config: PrismaticConfig) -> None:
613
+ super().__init__(config)
614
+
615
+ # [Validation] Lightweight Validate on `config` Fields + Dependency Versions
616
+ if config.use_fused_vision_backbone is None:
617
+ raise ValueError("Missing config field `use_fused_vision_backbone`")
618
+
619
+ if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
620
+ raise NotImplementedError(
621
+ "TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
622
+ "if you urgently need support for latest TIMM versions."
623
+ )
624
+
625
+ if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
626
+ logger.warning(
627
+ f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
628
+ f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
629
+ f"there might be inference-time regressions due to dependency changes. If in doubt, please"
630
+ f"use the above versions."
631
+ )
632
+
633
+ # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
634
+ self.vision_backbone = PrismaticVisionBackbone(
635
+ config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
636
+ )
637
+
638
+ # Create Multimodal Projector
639
+ self.projector = PrismaticProjector(
640
+ config.use_fused_vision_backbone,
641
+ vision_dim=self.vision_backbone.embed_dim,
642
+ llm_dim=config.text_config.hidden_size,
643
+ )
644
+
645
+ # Instantiate LLM Backbone
646
+ self.language_model = AutoModelForCausalLM.from_config(
647
+ config.text_config, attn_implementation=config._attn_implementation
648
+ )
649
+ self.vocab_size = config.text_config.vocab_size
650
+ self.pad_token_id = config.pad_token_id
651
+ self.llm_dim = config.text_config.hidden_size
652
+ self.use_text_token_gate = config.use_text_token_gate
653
+ self.text_token_gate_use_text_summary = getattr(config, "text_token_gate_use_text_summary", True)
654
+ self.text_token_gate_use_vision_tokens = getattr(config, "text_token_gate_use_vision_tokens", True)
655
+ self.text_token_gate_budget = config.text_token_gate_budget
656
+ self.text_token_gate_budget_loss_weight = config.text_token_gate_budget_loss_weight
657
+ self.text_token_gate_linear_mean_penalty_weight = config.text_token_gate_linear_mean_penalty_weight
658
+ self.text_token_gate_pool_instruction_only = getattr(config, "text_token_gate_pool_instruction_only", False)
659
+ self.text_token_gate_hidden_dim = config.text_token_gate_hidden_dim
660
+ self.text_token_gate_text_pooling_mode = config.text_token_gate_text_pooling_mode
661
+ self.text_token_gate_text_pool_hidden_dim = config.text_token_gate_text_pool_hidden_dim
662
+ self.text_token_gate_cross_attention_dim = config.text_token_gate_cross_attention_dim
663
+ self.text_token_gate_cross_attention_heads = getattr(config, "text_token_gate_cross_attention_heads", 1)
664
+ self.text_token_gate_mlp_depth = getattr(config, "text_token_gate_mlp_depth", 1)
665
+ self.contrastive_visual_tau = getattr(config, "contrastive_visual_tau", 0.1)
666
+ self.contrastive_text_tau = getattr(config, "contrastive_text_tau", 1.0)
667
+ self.gazing_mode = getattr(config, "gazing_mode", "mlp")
668
+ self.self_attn_dim = getattr(config, "self_attn_dim", 512)
669
+ self.self_attn_heads = getattr(config, "self_attn_heads", 8)
670
+ self.self_attn_layers = getattr(config, "self_attn_layers", 1)
671
+ self.layer_gate_mode = getattr(config, "layer_gate_mode", "none")
672
+ self.layer_gate_threshold = getattr(config, "layer_gate_threshold", 0.15)
673
+ self.layer_gate_strength = getattr(config, "layer_gate_strength", 0.5)
674
+ self.text_token_gate = None
675
+ if self.use_text_token_gate:
676
+ self.text_token_gate = TextConditionedTokenGate(
677
+ llm_dim=config.text_config.hidden_size,
678
+ use_text_summary=self.text_token_gate_use_text_summary,
679
+ use_vision_tokens=self.text_token_gate_use_vision_tokens,
680
+ hidden_dim=self.text_token_gate_hidden_dim,
681
+ text_pooling_mode=self.text_token_gate_text_pooling_mode,
682
+ text_pool_hidden_dim=self.text_token_gate_text_pool_hidden_dim,
683
+ cross_attention_dim=self.text_token_gate_cross_attention_dim,
684
+ cross_attention_heads=self.text_token_gate_cross_attention_heads,
685
+ gate_mlp_depth=self.text_token_gate_mlp_depth,
686
+ contrastive_visual_tau=self.contrastive_visual_tau,
687
+ contrastive_text_tau=self.contrastive_text_tau,
688
+ gazing_mode=self.gazing_mode,
689
+ self_attn_dim=self.self_attn_dim,
690
+ self_attn_heads=self.self_attn_heads,
691
+ self_attn_layers=self.self_attn_layers,
692
+ )
693
+ self._last_token_gate_budget_loss = None
694
+ self._last_token_gate_linear_mean_penalty = None
695
+ self._last_token_gate_regularization_loss = None
696
+ self._debug_last_token_gate = None
697
+ self._debug_last_token_gate_mean = None
698
+ self._debug_last_token_gate_budget_loss = None
699
+ self._debug_last_token_gate_linear_mean_penalty = None
700
+ self._debug_last_text_pool_weights = None
701
+ self._debug_last_text_pool_mask = None
702
+ self._debug_last_contrastive_scores = None
703
+
704
+ # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
705
+ self.post_init()
706
+
707
+ # === `PreTrainedModel` Boilerplate ===
708
+ def get_input_embeddings(self) -> nn.Module:
709
+ return self.language_model.get_input_embeddings()
710
+
711
+ def set_input_embeddings(self, value: nn.Module) -> None:
712
+ self.language_model.set_input_embeddings(value)
713
+
714
+ def get_output_embeddings(self) -> nn.Module:
715
+ return self.language_model.get_output_embeddings()
716
+
717
+ def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
718
+ self.language_model.set_output_embeddings(new_embeddings)
719
+
720
+ def get_decoder(self) -> nn.Module:
721
+ return self.language_model.get_decoder()
722
+
723
+ def set_decoder(self, decoder: nn.Module) -> None:
724
+ self.language_model.set_decoder(decoder)
725
+
726
+ def tie_weights(self) -> None:
727
+ self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
728
+
729
+ def resize_token_embeddings(
730
+ self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
731
+ ) -> nn.Embedding:
732
+ updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
733
+
734
+ # Update config/instance variables
735
+ self.config.text_config.vocab_size = updated_embeddings.num_embeddings
736
+ self.vocab_size = updated_embeddings.num_embeddings
737
+
738
+ return updated_embeddings
739
+
740
+ def _configure_text_token_gate_from_config(self) -> None:
741
+ self.use_text_token_gate = self.config.use_text_token_gate
742
+ self.text_token_gate_use_text_summary = getattr(self.config, "text_token_gate_use_text_summary", True)
743
+ self.text_token_gate_use_vision_tokens = getattr(self.config, "text_token_gate_use_vision_tokens", True)
744
+ self.text_token_gate_budget = self.config.text_token_gate_budget
745
+ self.text_token_gate_budget_loss_weight = self.config.text_token_gate_budget_loss_weight
746
+ self.text_token_gate_linear_mean_penalty_weight = self.config.text_token_gate_linear_mean_penalty_weight
747
+ self.text_token_gate_pool_instruction_only = getattr(self.config, "text_token_gate_pool_instruction_only", False)
748
+ self.text_token_gate_hidden_dim = self.config.text_token_gate_hidden_dim
749
+ self.text_token_gate_text_pooling_mode = self.config.text_token_gate_text_pooling_mode
750
+ self.text_token_gate_text_pool_hidden_dim = self.config.text_token_gate_text_pool_hidden_dim
751
+ self.text_token_gate_cross_attention_dim = self.config.text_token_gate_cross_attention_dim
752
+ self.text_token_gate_cross_attention_heads = getattr(self.config, "text_token_gate_cross_attention_heads", 1)
753
+ self.text_token_gate_mlp_depth = getattr(self.config, "text_token_gate_mlp_depth", 1)
754
+ self.contrastive_visual_tau = getattr(self.config, "contrastive_visual_tau", 0.1)
755
+ self.contrastive_text_tau = getattr(self.config, "contrastive_text_tau", 1.0)
756
+ self.gazing_mode = getattr(self.config, "gazing_mode", "mlp")
757
+ self.self_attn_dim = getattr(self.config, "self_attn_dim", 512)
758
+ self.self_attn_heads = getattr(self.config, "self_attn_heads", 8)
759
+ self.self_attn_layers = getattr(self.config, "self_attn_layers", 1)
760
+ self.layer_gate_mode = getattr(self.config, "layer_gate_mode", "none")
761
+ self.layer_gate_threshold = getattr(self.config, "layer_gate_threshold", 0.15)
762
+ self.layer_gate_strength = getattr(self.config, "layer_gate_strength", 0.5)
763
+
764
+ if not self.use_text_token_gate:
765
+ self.text_token_gate = None
766
+ self._last_token_gate_budget_loss = None
767
+ self._last_token_gate_linear_mean_penalty = None
768
+ self._last_token_gate_regularization_loss = None
769
+ self._debug_last_token_gate = None
770
+ self._debug_last_token_gate_mean = None
771
+ self._debug_last_token_gate_budget_loss = None
772
+ self._debug_last_token_gate_linear_mean_penalty = None
773
+ self._debug_last_text_pool_weights = None
774
+ self._debug_last_text_pool_mask = None
775
+ self._debug_last_contrastive_scores = None
776
+ return
777
+
778
+ needs_new_gate = self.text_token_gate is None
779
+ if not needs_new_gate:
780
+ gate_module = self.text_token_gate
781
+ if hasattr(gate_module, "original_module"):
782
+ gate_module = gate_module.original_module
783
+ elif hasattr(gate_module, "modules_to_save") and len(gate_module.modules_to_save) > 0:
784
+ gate_module = next(iter(gate_module.modules_to_save.values()))
785
+
786
+ gate_hidden_dim = getattr(gate_module, "hidden_dim", None)
787
+ if gate_hidden_dim is None and getattr(gate_module, "gate", None) is not None:
788
+ gate_hidden_dim = gate_module.gate[0].out_features
789
+ gate_text_pooling_mode = getattr(gate_module, "text_pooling_mode", None)
790
+ gate_text_pool_hidden_dim = getattr(gate_module, "text_pool_hidden_dim", None)
791
+ gate_cross_attention_dim = getattr(gate_module, "cross_attention_dim", None)
792
+ gate_cross_attention_heads = getattr(gate_module, "cross_attention_heads", 1)
793
+ gate_mlp_depth = getattr(gate_module, "gate_mlp_depth", 1)
794
+ gate_contrastive_visual_tau = getattr(gate_module, "contrastive_visual_tau", 0.1)
795
+ gate_contrastive_text_tau = getattr(gate_module, "contrastive_text_tau", 1.0)
796
+ gate_gazing_mode = getattr(gate_module, "gazing_mode", "mlp")
797
+ gate_use_text_summary = getattr(gate_module, "use_text_summary", True)
798
+ gate_use_vision_tokens = getattr(gate_module, "use_vision_tokens", True)
799
+ gate_self_attn_dim = getattr(gate_module, "self_attn_dim", 512)
800
+ gate_self_attn_heads = getattr(gate_module, "self_attn_heads", 8)
801
+ gate_self_attn_layers = getattr(gate_module, "self_attn_layers", 1)
802
+ needs_new_gate = (
803
+ gate_use_text_summary != self.text_token_gate_use_text_summary
804
+ or gate_use_vision_tokens != self.text_token_gate_use_vision_tokens
805
+ or
806
+ gate_hidden_dim != self.text_token_gate_hidden_dim
807
+ or gate_text_pooling_mode != self.text_token_gate_text_pooling_mode
808
+ or gate_mlp_depth != self.text_token_gate_mlp_depth
809
+ or gate_gazing_mode != self.gazing_mode
810
+ or (
811
+ self.text_token_gate_text_pooling_mode == "mlp"
812
+ and gate_text_pool_hidden_dim != self.text_token_gate_text_pool_hidden_dim
813
+ )
814
+ or (
815
+ self.text_token_gate_text_pooling_mode == "cross_attention"
816
+ and gate_cross_attention_dim != self.text_token_gate_cross_attention_dim
817
+ )
818
+ or (
819
+ self.text_token_gate_text_pooling_mode == "cross_attention"
820
+ and gate_cross_attention_heads != self.text_token_gate_cross_attention_heads
821
+ )
822
+ or (
823
+ self.text_token_gate_text_pooling_mode == "contrastive_alignment_score"
824
+ and gate_contrastive_visual_tau != self.contrastive_visual_tau
825
+ )
826
+ or (
827
+ self.text_token_gate_text_pooling_mode == "contrastive_alignment_score"
828
+ and gate_contrastive_text_tau != self.contrastive_text_tau
829
+ )
830
+ or (self.gazing_mode == "self_attention" and gate_self_attn_dim != self.self_attn_dim)
831
+ or (self.gazing_mode == "self_attention" and gate_self_attn_heads != self.self_attn_heads)
832
+ or (self.gazing_mode == "self_attention" and gate_self_attn_layers != self.self_attn_layers)
833
+ )
834
+
835
+ if needs_new_gate:
836
+ self.text_token_gate = TextConditionedTokenGate(
837
+ llm_dim=self.llm_dim,
838
+ use_text_summary=self.text_token_gate_use_text_summary,
839
+ use_vision_tokens=self.text_token_gate_use_vision_tokens,
840
+ hidden_dim=self.text_token_gate_hidden_dim,
841
+ text_pooling_mode=self.text_token_gate_text_pooling_mode,
842
+ text_pool_hidden_dim=self.text_token_gate_text_pool_hidden_dim,
843
+ cross_attention_dim=self.text_token_gate_cross_attention_dim,
844
+ cross_attention_heads=self.text_token_gate_cross_attention_heads,
845
+ gate_mlp_depth=self.text_token_gate_mlp_depth,
846
+ contrastive_visual_tau=self.contrastive_visual_tau,
847
+ contrastive_text_tau=self.contrastive_text_tau,
848
+ gazing_mode=self.gazing_mode,
849
+ self_attn_dim=self.self_attn_dim,
850
+ self_attn_heads=self.self_attn_heads,
851
+ self_attn_layers=self.self_attn_layers,
852
+ )
853
+ self.text_token_gate.apply(self._init_weights)
854
+
855
+ reference_param = next(self.projector.parameters())
856
+ self.text_token_gate = self.text_token_gate.to(
857
+ device=reference_param.device,
858
+ dtype=reference_param.dtype,
859
+ )
860
+
861
+ def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
862
+ """
863
+ Replace embeddings in input_embeddings at positions where all_actions_mask is True
864
+ with embeddings from noisy_action_features, using vectorized operations.
865
+
866
+ Args:
867
+ input_embeddings: Tensor of shape (B, S, D)
868
+ all_actions_mask: Boolean tensor of shape (B, S)
869
+ noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
870
+
871
+ Returns:
872
+ Modified input_embeddings tensor
873
+ """
874
+ # Clone input to avoid modifying the original tensor
875
+ new_input_embeddings = input_embeddings.clone()
876
+
877
+ # Create a tensor with the same shape of input_embeddings to hold the noisy action features
878
+ repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
879
+
880
+ # Create batch indices for splicing
881
+ batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
882
+ batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
883
+
884
+ # Get indices where mask is True for each sample
885
+ masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
886
+
887
+ # Move the noisy action features into their correct positions
888
+ repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
889
+
890
+ # Combine original input embeddings and noisy action embeddings using the mask
891
+ new_input_embeddings = torch.where(
892
+ all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
893
+ )
894
+
895
+ return new_input_embeddings
896
+
897
+ def _process_action_masks(self, labels):
898
+ """Helper to get action masks from labels"""
899
+ current_action_mask = get_current_action_mask(labels)
900
+ next_actions_mask = get_next_actions_mask(labels)
901
+ all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
902
+ return all_actions_mask
903
+
904
+ def _process_vision_features(self, pixel_values, language_embeddings=None, use_film=False):
905
+ """Process vision features with optional FiLM conditioning"""
906
+ if use_film:
907
+ # FiLM: Infuse language inputs into visual features
908
+ patch_features = self.vision_backbone(pixel_values, language_embeddings) # (bsz, 256 * num_images, D)
909
+ else:
910
+ patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
911
+
912
+ # Project patch embeddings into language embedding space
913
+ return self.projector(patch_features)
914
+
915
+ def _apply_text_token_gate(
916
+ self,
917
+ projected_patch_embeddings,
918
+ language_embeddings,
919
+ language_attention_mask=None,
920
+ language_text_pool_mask=None,
921
+ ):
922
+ patches_per_image = self.vision_backbone.get_num_patches()
923
+ gated_tokens, token_gate, text_pool_weights = self.text_token_gate(
924
+ projected_patch_embeddings,
925
+ language_embeddings,
926
+ language_attention_mask,
927
+ language_text_pool_mask,
928
+ patches_per_image=patches_per_image,
929
+ )
930
+ gate_mean_per_sample = token_gate.mean(dim=1).squeeze(-1) # [B], global mean for logging
931
+ if patches_per_image > 0 and token_gate.shape[1] % patches_per_image == 0:
932
+ num_images = token_gate.shape[1] // patches_per_image
933
+ gate_mean_per_image = token_gate.reshape(
934
+ token_gate.shape[0], num_images, patches_per_image, token_gate.shape[-1]
935
+ ).mean(dim=2).squeeze(-1) # [B, num_images]
936
+ budget_loss = ((gate_mean_per_image - self.text_token_gate_budget) ** 2).mean()
937
+ linear_mean_penalty = gate_mean_per_image.mean()
938
+ else:
939
+ logger.warning(
940
+ "Could not split token gate by image for per-camera budget loss; "
941
+ f"tokens={token_gate.shape[1]}, patches_per_image={patches_per_image}. "
942
+ "Falling back to global token-gate mean."
943
+ )
944
+ budget_loss = ((gate_mean_per_sample - self.text_token_gate_budget) ** 2).mean()
945
+ linear_mean_penalty = gate_mean_per_sample.mean()
946
+ regularization_terms = []
947
+ if self.text_token_gate_budget_loss_weight != 0:
948
+ regularization_terms.append(self.text_token_gate_budget_loss_weight * budget_loss)
949
+ if self.text_token_gate_linear_mean_penalty_weight != 0:
950
+ regularization_terms.append(self.text_token_gate_linear_mean_penalty_weight * linear_mean_penalty)
951
+ regularization_loss = None
952
+ if regularization_terms:
953
+ regularization_loss = sum(regularization_terms)
954
+ self._last_token_gate_budget_loss = budget_loss
955
+ self._last_token_gate_linear_mean_penalty = linear_mean_penalty
956
+ self._last_token_gate_regularization_loss = regularization_loss
957
+ self._debug_last_token_gate = token_gate.detach()
958
+ self._debug_last_token_gate_mean = token_gate.mean().detach()
959
+ self._debug_last_token_gate_budget_loss = budget_loss.detach()
960
+ self._debug_last_token_gate_linear_mean_penalty = linear_mean_penalty.detach()
961
+ self._debug_last_text_pool_weights = text_pool_weights.detach() if text_pool_weights is not None else None
962
+ self._debug_last_text_pool_mask = getattr(self.text_token_gate, "last_text_pool_mask", None)
963
+ self._debug_last_contrastive_scores = getattr(self.text_token_gate, "last_contrastive_scores", None)
964
+ return gated_tokens, token_gate, regularization_loss
965
+
966
+ def _get_language_model_layers(self) -> Optional[nn.ModuleList]:
967
+ model = getattr(self.language_model, "model", self.language_model)
968
+ if hasattr(model, "layers"):
969
+ return model.layers
970
+ decoder = getattr(model, "decoder", None)
971
+ if decoder is not None and hasattr(decoder, "layers"):
972
+ return decoder.layers
973
+ return None
974
+
975
+ def _build_layer_gate(self, token_gate: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
976
+ if (
977
+ token_gate is None
978
+ or self.layer_gate_mode == "none"
979
+ or self.layer_gate_strength <= 0
980
+ or self.layer_gate_mode != "threshold"
981
+ ):
982
+ return None
983
+
984
+ gate = token_gate.clamp(min=0.0, max=1.0)
985
+ low_gate_mask = (gate < self.layer_gate_threshold).to(dtype=gate.dtype)
986
+ layer_gate = 1.0 - low_gate_mask * self.layer_gate_strength * (1.0 - gate)
987
+ return layer_gate
988
+
989
+ def _run_language_model_with_layer_gate(
990
+ self,
991
+ token_gate: Optional[torch.Tensor] = None,
992
+ visual_start_idx: int = 1,
993
+ **language_model_kwargs,
994
+ ):
995
+ layer_gate = self._build_layer_gate(token_gate)
996
+ if layer_gate is None:
997
+ return self.language_model(**language_model_kwargs)
998
+
999
+ layers = self._get_language_model_layers()
1000
+ if layers is None:
1001
+ logger.warning("Could not find language model decoder layers; skipping layer gate.")
1002
+ return self.language_model(**language_model_kwargs)
1003
+
1004
+ visual_token_count = layer_gate.shape[1]
1005
+
1006
+ def apply_to_hidden_states(hidden_states: torch.Tensor) -> torch.Tensor:
1007
+ visual_end_idx = visual_start_idx + visual_token_count
1008
+ if hidden_states.ndim != 3 or hidden_states.shape[1] < visual_end_idx:
1009
+ return hidden_states
1010
+ gate = layer_gate.to(device=hidden_states.device, dtype=hidden_states.dtype)
1011
+ return torch.cat(
1012
+ [
1013
+ hidden_states[:, :visual_start_idx, :],
1014
+ hidden_states[:, visual_start_idx:visual_end_idx, :] * gate,
1015
+ hidden_states[:, visual_end_idx:, :],
1016
+ ],
1017
+ dim=1,
1018
+ )
1019
+
1020
+ def pre_hook_with_kwargs(module, args, kwargs):
1021
+ if args and torch.is_tensor(args[0]):
1022
+ args = (apply_to_hidden_states(args[0]),) + args[1:]
1023
+ elif "hidden_states" in kwargs and torch.is_tensor(kwargs["hidden_states"]):
1024
+ kwargs = dict(kwargs)
1025
+ kwargs["hidden_states"] = apply_to_hidden_states(kwargs["hidden_states"])
1026
+ return args, kwargs
1027
+
1028
+ def pre_hook(module, args):
1029
+ if args and torch.is_tensor(args[0]):
1030
+ return (apply_to_hidden_states(args[0]),) + args[1:]
1031
+ return args
1032
+
1033
+ handles = []
1034
+ try:
1035
+ for layer in layers:
1036
+ try:
1037
+ handles.append(layer.register_forward_pre_hook(pre_hook_with_kwargs, with_kwargs=True))
1038
+ except TypeError:
1039
+ handles.append(layer.register_forward_pre_hook(pre_hook))
1040
+ return self.language_model(**language_model_kwargs)
1041
+ finally:
1042
+ for handle in handles:
1043
+ handle.remove()
1044
+
1045
+ def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
1046
+ """Process proprioceptive features and append to vision features"""
1047
+ if proprio_projector is not None and proprio is not None:
1048
+ # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
1049
+ # proprio: (bsz, proprio_dim) or (propro_dim,)
1050
+ proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
1051
+ proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
1052
+ proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
1053
+ # For simplicity, just append proprio token to the end of projected vision patch tokens
1054
+ return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
1055
+ return projected_patch_embeddings
1056
+
1057
+ def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
1058
+ """Build multimodal embeddings and attention mask"""
1059
+ # Update attention mask
1060
+ projected_patch_attention_mask = None
1061
+ if attention_mask is not None:
1062
+ projected_patch_attention_mask = torch.full(
1063
+ (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
1064
+ fill_value=True,
1065
+ dtype=attention_mask.dtype,
1066
+ device=attention_mask.device,
1067
+ )
1068
+
1069
+ # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
1070
+ multimodal_embeddings = torch.cat(
1071
+ [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
1072
+ )
1073
+
1074
+ multimodal_attention_mask = None
1075
+ if attention_mask is not None:
1076
+ multimodal_attention_mask = torch.cat(
1077
+ [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
1078
+ )
1079
+
1080
+ return multimodal_embeddings, multimodal_attention_mask
1081
+
1082
+ def _build_multimodal_labels(self, labels, projected_patch_embeddings):
1083
+ """Build multimodal labels with IGNORE_INDEX for patch embeddings"""
1084
+ if labels is not None:
1085
+ projected_patch_labels = torch.full(
1086
+ (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
1087
+ fill_value=IGNORE_INDEX,
1088
+ dtype=labels.dtype,
1089
+ device=labels.device,
1090
+ )
1091
+ return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1)
1092
+ return None
1093
+
1094
+ # === Core Prismatic VLM `forward()` Logic ===
1095
+ def forward(
1096
+ self,
1097
+ input_ids: Optional[torch.LongTensor] = None,
1098
+ attention_mask: Optional[torch.Tensor] = None,
1099
+ pixel_values: Optional[torch.FloatTensor] = None,
1100
+ labels: Optional[torch.LongTensor] = None,
1101
+ text_token_pool_mask: Optional[torch.Tensor] = None,
1102
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1103
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1104
+ use_cache: Optional[bool] = None,
1105
+ output_attentions: Optional[bool] = None,
1106
+ output_hidden_states: Optional[bool] = None,
1107
+ output_projector_features: Optional[bool] = None,
1108
+ return_dict: Optional[bool] = None,
1109
+ proprio=None,
1110
+ proprio_projector=None,
1111
+ noisy_actions=None,
1112
+ noisy_action_projector=None,
1113
+ diffusion_timestep_embeddings=None,
1114
+ use_film: bool = False,
1115
+ ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
1116
+ """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
1117
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1118
+ output_hidden_states = (
1119
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1120
+ )
1121
+ output_projector_features = output_projector_features if output_projector_features is not None else False
1122
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1123
+
1124
+ # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
1125
+ use_cache = use_cache and not self.training
1126
+
1127
+ # Clear token-gate stats so they never leak across forwards.
1128
+ self._last_token_gate_budget_loss = None
1129
+ self._last_token_gate_linear_mean_penalty = None
1130
+ self._last_token_gate_regularization_loss = None
1131
+ self._debug_last_token_gate = None
1132
+ self._debug_last_token_gate_mean = None
1133
+ self._debug_last_token_gate_budget_loss = None
1134
+ self._debug_last_token_gate_linear_mean_penalty = None
1135
+ self._debug_last_text_pool_weights = None
1136
+ self._debug_last_text_pool_mask = None
1137
+ self._debug_last_contrastive_scores = None
1138
+
1139
+ # Instantiate Placeholder for Projector Features
1140
+ projected_patch_embeddings = None
1141
+ token_gate = None
1142
+ token_gate_regularization_loss = None
1143
+
1144
+ # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
1145
+ if input_ids.shape[1] == 1:
1146
+ assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
1147
+ assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
1148
+ assert labels is None, "Unexpected key `labels` provided during cached generation!"
1149
+
1150
+ language_model_output = self.language_model(
1151
+ input_ids=input_ids,
1152
+ attention_mask=None,
1153
+ position_ids=None,
1154
+ past_key_values=past_key_values,
1155
+ inputs_embeds=None,
1156
+ labels=None,
1157
+ use_cache=use_cache,
1158
+ output_attentions=output_attentions,
1159
+ output_hidden_states=output_hidden_states,
1160
+ return_dict=return_dict,
1161
+ )
1162
+
1163
+ # === Handle Unimodal Forward ===
1164
+ elif pixel_values is None:
1165
+ assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
1166
+ assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
1167
+
1168
+ language_model_output = self.language_model(
1169
+ input_ids=input_ids,
1170
+ attention_mask=attention_mask,
1171
+ position_ids=None,
1172
+ past_key_values=None,
1173
+ inputs_embeds=None,
1174
+ labels=labels,
1175
+ use_cache=use_cache,
1176
+ output_attentions=output_attentions,
1177
+ output_hidden_states=output_hidden_states,
1178
+ return_dict=return_dict,
1179
+ )
1180
+
1181
+ # === Handle Multimodal Forward ===
1182
+ elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
1183
+ assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
1184
+
1185
+ # Get input embeddings (from language model embeddings)
1186
+ input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
1187
+
1188
+ # Extract action masks
1189
+ all_actions_mask = self._process_action_masks(labels)
1190
+
1191
+ # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
1192
+ language_embeddings = input_embeddings[~all_actions_mask].reshape(
1193
+ input_embeddings.shape[0], -1, input_embeddings.shape[2]
1194
+ ) # (B, lang_seq_len, llm_dim)
1195
+ language_attention_mask = None
1196
+ if attention_mask is not None:
1197
+ language_attention_mask = attention_mask[~all_actions_mask].reshape(input_embeddings.shape[0], -1)
1198
+ language_text_pool_mask = None
1199
+ if text_token_pool_mask is not None:
1200
+ text_token_pool_mask = text_token_pool_mask.to(device=input_embeddings.device, dtype=torch.bool)
1201
+ if text_token_pool_mask.shape != all_actions_mask.shape:
1202
+ logger.warning(
1203
+ "text_token_pool_mask shape does not match input/action mask shape; "
1204
+ f"got {text_token_pool_mask.shape} vs {all_actions_mask.shape}. Ignoring pool mask."
1205
+ )
1206
+ else:
1207
+ language_text_pool_mask = text_token_pool_mask[~all_actions_mask].reshape(
1208
+ input_embeddings.shape[0], -1
1209
+ )
1210
+
1211
+ # Get visual features
1212
+ projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
1213
+
1214
+ # Apply the gate only to projected visual tokens before proprio is appended.
1215
+ if self.use_text_token_gate:
1216
+ projected_patch_embeddings, token_gate, token_gate_regularization_loss = self._apply_text_token_gate(
1217
+ projected_patch_embeddings,
1218
+ language_embeddings,
1219
+ language_attention_mask,
1220
+ language_text_pool_mask,
1221
+ )
1222
+
1223
+ # Add proprioceptive state if provided
1224
+ projected_patch_embeddings = self._process_proprio_features(
1225
+ projected_patch_embeddings, proprio, proprio_projector
1226
+ )
1227
+
1228
+ # [Diffusion] Add diffusion timestep embedding if provided
1229
+ if diffusion_timestep_embeddings is not None:
1230
+ # For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
1231
+ projected_patch_embeddings = torch.cat(
1232
+ (projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
1233
+ )
1234
+
1235
+ # Process action embeddings
1236
+ if noisy_actions is not None:
1237
+ # Get mask corresponding to all action tokens
1238
+ all_actions_mask = self._process_action_masks(labels)
1239
+
1240
+ # Reshape noisy actions into individual action tokens
1241
+ # noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
1242
+ B = noisy_actions.shape[0]
1243
+ noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
1244
+
1245
+ # Project noisy action tokens into language model embedding space
1246
+ noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
1247
+
1248
+ # Replace embeddings of the action tokens with noisy action embeddings
1249
+ input_embeddings = self._replace_input_embeddings(
1250
+ input_embeddings, all_actions_mask, noisy_action_features
1251
+ )
1252
+ else:
1253
+ # Replace the embeddings of the action tokens with zeros
1254
+ # (Later on, the positional embeddings will be added to them)
1255
+ all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
1256
+ input_embeddings = input_embeddings * ~all_actions_mask
1257
+
1258
+ # Build multimodal embeddings & attention mask
1259
+ multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
1260
+ input_embeddings, projected_patch_embeddings, attention_mask
1261
+ )
1262
+
1263
+ # Build labels for multimodal sequence if needed
1264
+ multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
1265
+
1266
+ # Dispatch to language model
1267
+ language_model_output = self._run_language_model_with_layer_gate(
1268
+ token_gate=token_gate,
1269
+ input_ids=None,
1270
+ attention_mask=multimodal_attention_mask,
1271
+ position_ids=None,
1272
+ past_key_values=None,
1273
+ inputs_embeds=multimodal_embeddings,
1274
+ labels=multimodal_labels,
1275
+ use_cache=use_cache,
1276
+ output_attentions=output_attentions,
1277
+ output_hidden_states=output_hidden_states,
1278
+ return_dict=return_dict,
1279
+ )
1280
+
1281
+ # === Otherwise =>> Assume Invalid! ===
1282
+ elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
1283
+ raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
1284
+
1285
+ else:
1286
+ raise ValueError(
1287
+ "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
1288
+ f"=> `input_ids` = {input_ids is not None}\n"
1289
+ f"=> `attention_mask` = {attention_mask is not None}\n"
1290
+ f"=> `pixel_values` = {pixel_values is not None}\n"
1291
+ f"=> `labels` = {labels is not None}\n"
1292
+ f"=> `input_embeds` = {inputs_embeds is not None}\n"
1293
+ f"=> `past_key_values` = {past_key_values is not None}\n"
1294
+ f"=> `use_cache` = {use_cache}"
1295
+ )
1296
+
1297
+ total_loss = language_model_output.loss if return_dict else (language_model_output[0] if labels is not None else None)
1298
+ if token_gate_regularization_loss is not None and total_loss is not None:
1299
+ total_loss = total_loss + token_gate_regularization_loss
1300
+
1301
+ # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
1302
+ if not return_dict:
1303
+ if total_loss is not None:
1304
+ language_model_output = (total_loss,) + tuple(language_model_output[1:])
1305
+ if output_projector_features and (projected_patch_embeddings is not None):
1306
+ return *language_model_output, projected_patch_embeddings
1307
+
1308
+ return language_model_output
1309
+
1310
+ return PrismaticCausalLMOutputWithPast(
1311
+ loss=total_loss,
1312
+ logits=language_model_output.logits,
1313
+ past_key_values=language_model_output.past_key_values,
1314
+ hidden_states=language_model_output.hidden_states,
1315
+ attentions=language_model_output.attentions,
1316
+ projector_features=projected_patch_embeddings,
1317
+ )
1318
+
1319
+ # === GenerationMixin Methods ===
1320
+ def prepare_inputs_for_generation(
1321
+ self,
1322
+ input_ids: Optional[torch.Tensor] = None,
1323
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1324
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1325
+ pixel_values: Optional[torch.FloatTensor] = None,
1326
+ attention_mask: Optional[torch.Tensor] = None,
1327
+ **kwargs: str,
1328
+ ) -> Dict[str, torch.Tensor]:
1329
+ """Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
1330
+ if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
1331
+ (inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
1332
+ ):
1333
+ raise ValueError("Generation with batch size > 1 is not currently supported!")
1334
+
1335
+ # Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
1336
+ if past_key_values is not None:
1337
+ input_ids = input_ids[:, -1:]
1338
+
1339
+ # If `input_embeds` are passed, we only want to use them in the 1st generation step
1340
+ if inputs_embeds is not None and past_key_values is None:
1341
+ model_inputs = {"input_embeds": inputs_embeds}
1342
+ else:
1343
+ model_inputs = {"input_ids": input_ids}
1344
+
1345
+ # Make sure `pixel_values` are preserved in `model_inputs`
1346
+ model_inputs.update(
1347
+ {
1348
+ "attention_mask": attention_mask,
1349
+ "pixel_values": pixel_values,
1350
+ "past_key_values": past_key_values,
1351
+ "use_cache": kwargs.get("use_cache"),
1352
+ }
1353
+ )
1354
+
1355
+ return model_inputs
1356
+
1357
+ # Defer to Language Model (all handle this differently, with different return types)
1358
+ def _reorder_cache(self, *args, **kwargs) -> Any:
1359
+ return self.language_model._reorder_cache(*args, **kwargs)
1360
+
1361
+
1362
+ class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
1363
+ config_class: PretrainedConfig = OpenVLAConfig
1364
+
1365
+ def __init__(self, config: OpenVLAConfig) -> None:
1366
+ super().__init__(config)
1367
+ self.norm_stats = config.norm_stats
1368
+
1369
+ # Compute action bins
1370
+ self.bins = np.linspace(-1, 1, config.n_action_bins)
1371
+ self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
1372
+
1373
+ # Compute vocab size for de-tokenization -- revert added "multiple of"
1374
+ self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
1375
+
1376
+ def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
1377
+ """Prepares input for action prediction by adding necessary tokens"""
1378
+ # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
1379
+ placeholder_action_token_ids = (
1380
+ torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
1381
+ )
1382
+ input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
1383
+
1384
+ # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
1385
+ stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
1386
+ input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
1387
+
1388
+ # Extend the attention mask to fit the new shape of input
1389
+ # Note: Only batch size == 1 supported right now
1390
+ mask_extension = (
1391
+ torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
1392
+ .to(attention_mask.device)
1393
+ .to(attention_mask.dtype)
1394
+ )
1395
+ attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
1396
+
1397
+ return input_ids, attention_mask
1398
+
1399
+ def _prepare_labels_for_action_prediction(self, labels, input_ids):
1400
+ """Creates labels tensor for action prediction if not provided"""
1401
+ # Extend labels tensor with fake action labels
1402
+ ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
1403
+ labels_extension = (
1404
+ torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
1405
+ * ARBITRARY_ACTION_TOKEN_IDX
1406
+ )
1407
+ labels = torch.cat([labels, labels_extension], dim=-1)
1408
+
1409
+ # Replace last label token with stop token
1410
+ labels[:, -1] = STOP_INDEX
1411
+
1412
+ return labels
1413
+
1414
+ def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
1415
+ """Unnormalize actions using dataset statistics"""
1416
+ action_norm_stats = self.get_action_stats(unnorm_key)
1417
+
1418
+ if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
1419
+ mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
1420
+ action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
1421
+ elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
1422
+ mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
1423
+ action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
1424
+ else:
1425
+ raise ValueError("Unsupported action/proprio normalization type detected!")
1426
+
1427
+ actions = np.where(
1428
+ mask,
1429
+ 0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
1430
+ normalized_actions,
1431
+ )
1432
+
1433
+ return actions
1434
+
1435
+ def _run_diffusion_prediction(
1436
+ self,
1437
+ input_embeddings,
1438
+ all_actions_mask,
1439
+ noise,
1440
+ action_head,
1441
+ projected_patch_embeddings,
1442
+ labels,
1443
+ attention_mask,
1444
+ NUM_PATCHES,
1445
+ NUM_PROMPT_TOKENS,
1446
+ noisy_action_projector,
1447
+ token_gate=None,
1448
+ ):
1449
+ """Run diffusion-based action prediction"""
1450
+ # Clone embedding for reuse in each timestep
1451
+ orig_projected_patch_embeddings = projected_patch_embeddings.clone()
1452
+ curr_noisy_actions = noise
1453
+
1454
+ # Reverse diffusion: Iteratively denoise to generate action prediction
1455
+ for t in action_head.noise_scheduler.timesteps:
1456
+ # Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
1457
+ # embedding, and diffusion timestep embedding)
1458
+ timesteps = torch.Tensor([t]).to(labels.device)
1459
+ diffusion_timestep_embeddings = (
1460
+ action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
1461
+ ) # (B, llm_dim)
1462
+ diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
1463
+
1464
+ # [Diffusion] Replace the embeddings of the action tokens with noisy actions
1465
+ # (Later on, the positional embeddings will be added to them)
1466
+
1467
+ # For simplicity, append diffusion timestep embedding to the end of projected vision tokens
1468
+ projected_patch_embeddings = torch.cat(
1469
+ (orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
1470
+ )
1471
+
1472
+ # Reshape and project noisy actions into language embedding space
1473
+ B = curr_noisy_actions.shape[0]
1474
+ orig_curr_noisy_actions_shape = curr_noisy_actions.shape
1475
+ curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
1476
+ noisy_action_features = noisy_action_projector(curr_noisy_actions)
1477
+ curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
1478
+
1479
+ # Replace action token embeddings with noisy action embeddings
1480
+ input_embeddings = self._replace_input_embeddings(
1481
+ input_embeddings.clone(), all_actions_mask, noisy_action_features
1482
+ )
1483
+
1484
+ # Build multimodal embeddings and attention mask
1485
+ multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
1486
+ input_embeddings, projected_patch_embeddings, attention_mask
1487
+ )
1488
+
1489
+ # Forward pass through language model
1490
+ language_model_output = self._run_language_model_with_layer_gate(
1491
+ token_gate=token_gate,
1492
+ input_ids=None,
1493
+ attention_mask=multimodal_attention_mask,
1494
+ position_ids=None,
1495
+ past_key_values=None,
1496
+ inputs_embeds=multimodal_embeddings,
1497
+ labels=None,
1498
+ use_cache=None,
1499
+ output_attentions=False,
1500
+ output_hidden_states=True,
1501
+ return_dict=True,
1502
+ )
1503
+
1504
+ # Extract hidden states for action portion of response
1505
+ last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
1506
+ actions_hidden_states = last_hidden_states[
1507
+ :,
1508
+ NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
1509
+ :,
1510
+ ] # (B, act_chunk_len, D)
1511
+
1512
+ # Predict noise and update noisy actions: x_t -> x_{t-1}
1513
+ noise_pred = action_head.predict_noise(actions_hidden_states)
1514
+ curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
1515
+
1516
+ curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
1517
+
1518
+ # Return final actions
1519
+ return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
1520
+
1521
+ def _regression_or_discrete_prediction(
1522
+ self,
1523
+ input_embeddings,
1524
+ all_actions_mask,
1525
+ projected_patch_embeddings,
1526
+ attention_mask,
1527
+ labels,
1528
+ NUM_PATCHES,
1529
+ NUM_PROMPT_TOKENS,
1530
+ action_head=None,
1531
+ token_gate=None,
1532
+ ):
1533
+ """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
1534
+ # Zero out action token embeddings
1535
+ all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
1536
+ input_embeddings = input_embeddings * ~all_actions_mask
1537
+
1538
+ # Build multimodal embeddings and attention mask
1539
+ multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
1540
+ input_embeddings, projected_patch_embeddings, attention_mask
1541
+ )
1542
+
1543
+ # Forward pass through language model
1544
+ language_model_output = self._run_language_model_with_layer_gate(
1545
+ token_gate=token_gate,
1546
+ input_ids=None,
1547
+ attention_mask=multimodal_attention_mask,
1548
+ position_ids=None,
1549
+ past_key_values=None,
1550
+ inputs_embeds=multimodal_embeddings,
1551
+ labels=None,
1552
+ use_cache=None,
1553
+ output_attentions=False,
1554
+ output_hidden_states=True,
1555
+ return_dict=True,
1556
+ )
1557
+
1558
+ # Extract hidden states for action tokens
1559
+ last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
1560
+ actions_hidden_states = last_hidden_states[
1561
+ :,
1562
+ NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
1563
+ :,
1564
+ ] # (B, act_chunk_len, D)
1565
+
1566
+ # Handle different prediction methods
1567
+ if action_head is not None:
1568
+ # L1 regression prediction
1569
+ normalized_actions = action_head.predict_action(actions_hidden_states)
1570
+ normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
1571
+ normalized_actions = normalized_actions.float().cpu().detach().numpy()
1572
+ else:
1573
+ # Discrete token-based prediction
1574
+ predicted_action_token_ids = (
1575
+ language_model_output.logits[
1576
+ :,
1577
+ NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
1578
+ ]
1579
+ .argmax(dim=2)
1580
+ .cpu()
1581
+ .numpy()
1582
+ )
1583
+ discretized_actions = self.vocab_size - predicted_action_token_ids
1584
+ discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
1585
+ normalized_actions = self.bin_centers[discretized_actions]
1586
+ normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
1587
+
1588
+ return normalized_actions, actions_hidden_states
1589
+
1590
+ def predict_action(
1591
+ self,
1592
+ input_ids: Optional[torch.LongTensor] = None,
1593
+ unnorm_key: Optional[str] = None,
1594
+ proprio=None,
1595
+ proprio_projector=None,
1596
+ action_head=None,
1597
+ noisy_action_projector=None,
1598
+ use_film: bool = False,
1599
+ text_token_pool_mask: Optional[torch.Tensor] = None,
1600
+ **kwargs: str,
1601
+ ) -> np.ndarray:
1602
+ """Predict actions from input sequence, with options for different prediction methods.
1603
+
1604
+ Args:
1605
+ input_ids: Input token ids
1606
+ unnorm_key: Key for unnormalization statistics
1607
+ proprio: Proprioceptive features
1608
+ proprio_projector: Projector for proprioceptive features
1609
+ action_head: Optional head for L1 regression or diffusion-based prediction
1610
+ noisy_action_projector: Projector for noisy actions in diffusion-based prediction
1611
+ use_film: Whether to use FiLM conditioning
1612
+ **kwargs: Additional arguments including pixel_values and attention_mask
1613
+
1614
+ Returns:
1615
+ Tuple of (unnormalized_actions, action_hidden_states)
1616
+ """
1617
+ self._last_token_gate_budget_loss = None
1618
+ self._last_token_gate_linear_mean_penalty = None
1619
+ self._last_token_gate_regularization_loss = None
1620
+ self._debug_last_token_gate = None
1621
+ self._debug_last_token_gate_mean = None
1622
+ self._debug_last_token_gate_budget_loss = None
1623
+ self._debug_last_token_gate_linear_mean_penalty = None
1624
+ self._debug_last_text_pool_weights = None
1625
+ self._debug_last_text_pool_mask = None
1626
+ self._debug_last_contrastive_scores = None
1627
+
1628
+ # If the special empty token ('') does not already appear after the colon (':') token in the prompt
1629
+ # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
1630
+ if text_token_pool_mask is not None:
1631
+ text_token_pool_mask = text_token_pool_mask.to(device=input_ids.device, dtype=torch.bool)
1632
+ if not torch.all(input_ids[:, -1] == 29871):
1633
+ input_ids = torch.cat(
1634
+ (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
1635
+ )
1636
+ if text_token_pool_mask is not None:
1637
+ text_token_pool_mask = torch.cat(
1638
+ [
1639
+ text_token_pool_mask,
1640
+ torch.zeros((text_token_pool_mask.shape[0], 1), device=input_ids.device, dtype=torch.bool),
1641
+ ],
1642
+ dim=1,
1643
+ )
1644
+
1645
+ pixel_values = kwargs["pixel_values"]
1646
+ attention_mask = kwargs["attention_mask"]
1647
+
1648
+ # Create fake labels tensor (needed for action mask)
1649
+ labels = input_ids.clone()
1650
+ labels[:] = IGNORE_INDEX
1651
+
1652
+ # Get number of tokens in prompt (excluding the start token)
1653
+ NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
1654
+
1655
+ # Prepare inputs by adding necessary tokens
1656
+ input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
1657
+ if text_token_pool_mask is not None:
1658
+ mask_extension_len = input_ids.shape[-1] - text_token_pool_mask.shape[-1]
1659
+ if mask_extension_len > 0:
1660
+ text_token_pool_mask = torch.cat(
1661
+ [
1662
+ text_token_pool_mask,
1663
+ torch.zeros(
1664
+ (text_token_pool_mask.shape[0], mask_extension_len),
1665
+ device=input_ids.device,
1666
+ dtype=torch.bool,
1667
+ ),
1668
+ ],
1669
+ dim=1,
1670
+ )
1671
+
1672
+ # Update labels tensor for action mask computation later
1673
+ labels = self._prepare_labels_for_action_prediction(labels, input_ids)
1674
+
1675
+ # Get input embeddings and action masks
1676
+ input_embeddings = self.get_input_embeddings()(input_ids)
1677
+ all_actions_mask = self._process_action_masks(labels)
1678
+
1679
+ # Extract language embeddings
1680
+ language_embeddings = input_embeddings[~all_actions_mask].reshape(
1681
+ input_embeddings.shape[0], -1, input_embeddings.shape[2]
1682
+ )
1683
+ language_attention_mask = None
1684
+ if attention_mask is not None:
1685
+ language_attention_mask = attention_mask[~all_actions_mask].reshape(input_embeddings.shape[0], -1)
1686
+ language_text_pool_mask = None
1687
+ if text_token_pool_mask is not None:
1688
+ if text_token_pool_mask.shape != all_actions_mask.shape:
1689
+ logger.warning(
1690
+ "text_token_pool_mask shape does not match eval input/action mask shape; "
1691
+ f"got {text_token_pool_mask.shape} vs {all_actions_mask.shape}. Ignoring pool mask."
1692
+ )
1693
+ else:
1694
+ language_text_pool_mask = text_token_pool_mask[~all_actions_mask].reshape(input_embeddings.shape[0], -1)
1695
+
1696
+ # Process vision features
1697
+ projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
1698
+
1699
+ # Apply the gate only to projected visual tokens before proprio is appended.
1700
+ token_gate = None
1701
+ if self.use_text_token_gate:
1702
+ projected_patch_embeddings, token_gate, _ = self._apply_text_token_gate(
1703
+ projected_patch_embeddings, language_embeddings, language_attention_mask, language_text_pool_mask
1704
+ )
1705
+
1706
+ # Add proprioceptive features if provided
1707
+ use_proprio = proprio_projector is not None and proprio is not None
1708
+ if use_proprio:
1709
+ proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
1710
+ projected_patch_embeddings = self._process_proprio_features(
1711
+ projected_patch_embeddings, proprio, proprio_projector
1712
+ )
1713
+
1714
+ # Use diffusion if provided, otherwise use regression or discrete prediction
1715
+ use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
1716
+
1717
+ # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
1718
+ NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
1719
+ if use_proprio:
1720
+ NUM_PATCHES += 1
1721
+ if use_diffusion:
1722
+ NUM_PATCHES += 1
1723
+
1724
+ if use_diffusion:
1725
+ # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion
1726
+ noise = torch.randn(
1727
+ size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype
1728
+ )
1729
+
1730
+ # Run diffusion-based prediction
1731
+ normalized_actions, actions_hidden_states = self._run_diffusion_prediction(
1732
+ input_embeddings,
1733
+ all_actions_mask,
1734
+ noise,
1735
+ action_head,
1736
+ projected_patch_embeddings,
1737
+ labels,
1738
+ attention_mask,
1739
+ NUM_PATCHES,
1740
+ NUM_PROMPT_TOKENS,
1741
+ noisy_action_projector,
1742
+ token_gate=token_gate,
1743
+ )
1744
+ else:
1745
+ # Run regression or discrete token-based prediction
1746
+ normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
1747
+ input_embeddings,
1748
+ all_actions_mask,
1749
+ projected_patch_embeddings,
1750
+ attention_mask,
1751
+ labels,
1752
+ NUM_PATCHES,
1753
+ NUM_PROMPT_TOKENS,
1754
+ action_head,
1755
+ token_gate=token_gate,
1756
+ )
1757
+
1758
+ # Unnormalize predicted actions
1759
+ actions = self._unnormalize_actions(normalized_actions, unnorm_key)
1760
+
1761
+ return actions, actions_hidden_states
1762
+
1763
+ @staticmethod
1764
+ def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
1765
+ """Validate and resolve the unnormalization key for action statistics"""
1766
+ if unnorm_key is None:
1767
+ assert len(norm_stats) == 1, (
1768
+ f"Your model was trained on more than one dataset, "
1769
+ f"please pass a `unnorm_key` from the following options to choose the statistics "
1770
+ f"used for un-normalizing actions: {norm_stats.keys()}"
1771
+ )
1772
+ unnorm_key = next(iter(norm_stats.keys()))
1773
+
1774
+ assert unnorm_key in norm_stats, (
1775
+ f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
1776
+ f"please choose from: {norm_stats.keys()}"
1777
+ )
1778
+ return unnorm_key
1779
+
1780
+ def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
1781
+ """Get the dimensionality of the policy's action space."""
1782
+ unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
1783
+ return len(self.norm_stats[unnorm_key]["action"]["min"])
1784
+
1785
+ def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
1786
+ """Get all the logged statistics for the given dataset."""
1787
+ unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
1788
+ return self.norm_stats[unnorm_key]["action"]
oft_training_config.json ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "dataset_name": "utokyo_xarm_pick_and_place_converted_externally_to_rlds",
3
+ "vla_path": "/workspace/kaixi/RealWorld/openvla_oft_runs/merged_public_checkpoints/oft_setting2_chunksize25_batch32_20k",
4
+ "use_film": true,
5
+ "num_images_in_input": 2,
6
+ "use_proprio": true,
7
+ "use_l1_regression": true,
8
+ "use_diffusion": false,
9
+ "use_lora": true,
10
+ "lora_rank": 32,
11
+ "num_actions_chunk": 25,
12
+ "action_dim": 7,
13
+ "proprio_dim": 6,
14
+ "action_proprio_normalization_type": "NormalizationType.BOUNDS_Q99"
15
+ }
preprocessor_config.json ADDED
@@ -0,0 +1,114 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoImageProcessor": "processing_prismatic.PrismaticImageProcessor",
4
+ "AutoProcessor": "processing_prismatic.PrismaticProcessor"
5
+ },
6
+ "image_processor_type": "PrismaticImageProcessor",
7
+ "image_resize_strategy": "resize-naive",
8
+ "input_sizes": [
9
+ [
10
+ 3,
11
+ 224,
12
+ 224
13
+ ],
14
+ [
15
+ 3,
16
+ 224,
17
+ 224
18
+ ]
19
+ ],
20
+ "interpolations": [
21
+ "bicubic",
22
+ "bicubic"
23
+ ],
24
+ "means": [
25
+ [
26
+ 0.485,
27
+ 0.456,
28
+ 0.406
29
+ ],
30
+ [
31
+ 0.5,
32
+ 0.5,
33
+ 0.5
34
+ ]
35
+ ],
36
+ "processor_class": "PrismaticProcessor",
37
+ "stds": [
38
+ [
39
+ 0.229,
40
+ 0.224,
41
+ 0.225
42
+ ],
43
+ [
44
+ 0.5,
45
+ 0.5,
46
+ 0.5
47
+ ]
48
+ ],
49
+ "tvf_crop_params": [
50
+ {
51
+ "output_size": [
52
+ 224,
53
+ 224
54
+ ]
55
+ },
56
+ {
57
+ "output_size": [
58
+ 224,
59
+ 224
60
+ ]
61
+ }
62
+ ],
63
+ "tvf_do_letterbox": false,
64
+ "tvf_letterbox_fill": null,
65
+ "tvf_normalize_params": [
66
+ {
67
+ "inplace": false,
68
+ "mean": [
69
+ 0.484375,
70
+ 0.455078125,
71
+ 0.40625
72
+ ],
73
+ "std": [
74
+ 0.228515625,
75
+ 0.2236328125,
76
+ 0.224609375
77
+ ]
78
+ },
79
+ {
80
+ "inplace": false,
81
+ "mean": [
82
+ 0.5,
83
+ 0.5,
84
+ 0.5
85
+ ],
86
+ "std": [
87
+ 0.5,
88
+ 0.5,
89
+ 0.5
90
+ ]
91
+ }
92
+ ],
93
+ "tvf_resize_params": [
94
+ {
95
+ "antialias": true,
96
+ "interpolation": 3,
97
+ "max_size": null,
98
+ "size": [
99
+ 224,
100
+ 224
101
+ ]
102
+ },
103
+ {
104
+ "antialias": true,
105
+ "interpolation": 3,
106
+ "max_size": null,
107
+ "size": [
108
+ 224,
109
+ 224
110
+ ]
111
+ }
112
+ ],
113
+ "use_fused_vision_backbone": true
114
+ }
processing_prismatic.py ADDED
@@ -0,0 +1,257 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ processing_prismatic.py
3
+
4
+ HuggingFace-style preprocessor definitions for Prismatic VLMs, inheriting from `ProcessorMixin`. Default configuration
5
+ specifies `siglip-224px+7b`.
6
+ """
7
+
8
+ from typing import Any, ClassVar, List, Optional, Tuple, Union
9
+
10
+ import timm.data
11
+ import torch
12
+ import torchvision.transforms.functional as TVF
13
+ from PIL import Image
14
+ from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
15
+ from transformers import PreTrainedTokenizerBase
16
+ from transformers.image_processing_utils import BatchFeature, ImageProcessingMixin
17
+ from transformers.processing_utils import ProcessorMixin
18
+ from transformers.tokenization_utils import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
19
+ from transformers.utils import TensorType
20
+
21
+
22
+ # === Image Processing ===
23
+ def letterbox_pad_transform(image: Image.Image, padding_fill_value: Tuple[int, int, int]) -> Image.Image:
24
+ """Given a PIL.Image, pad to square by adding a symmetric border around the height/width."""
25
+ (w, h), max_wh = image.size, max(image.size)
26
+ horizontal_pad, vertical_pad = int((max_wh - w) / 2), int((max_wh - h) / 2)
27
+ padding = (horizontal_pad, vertical_pad, horizontal_pad, vertical_pad)
28
+
29
+ return TVF.pad(image, padding, fill=padding_fill_value, padding_mode="constant")
30
+
31
+
32
+ class PrismaticImageProcessor(ImageProcessingMixin):
33
+ model_input_names: ClassVar[List[str]] = ["pixel_values"]
34
+
35
+ def __init__(
36
+ self,
37
+ use_fused_vision_backbone: bool = False,
38
+ image_resize_strategy: str = "letterbox",
39
+ input_sizes: Optional[List[Tuple[int, int, int]]] = None,
40
+ interpolations: Optional[List[str]] = None,
41
+ means: Optional[List[Tuple[float, float, float]]] = None,
42
+ stds: Optional[List[Tuple[float, float, float]]] = None,
43
+ **kwargs: str,
44
+ ) -> None:
45
+ """
46
+ Initialize a PrismaticImageProcessor as a wrapper around a torchvision transform; this transform will be
47
+ created by TIMM, and edited to follow our custom `image_resize_strategy` logic.
48
+
49
+ @param use_fused_vision_backbone: Boolean indicating single or fused (dual) vision backbone
50
+ @param image_resize_strategy: Prismatic image resize strategy in < resize-naive | resize-crop | letterbox >
51
+ @param input_size: [TIMM :: `data_cfg`] Input image size as tuple (channels, width, height)
52
+ @param interpolation: [TIMM :: `data_cfg`] Interpolation as string (default: "bicubic")
53
+ @param mean: [TIMM :: `data_cfg`] Normalization mean as float tuple (or two-tuple if `fused_backbone`)
54
+ @param std: [TIMM :: `data_cfg`] Normalization std as float tuple (or two-tuple if `fused_backbone`)
55
+ """
56
+ self.use_fused_vision_backbone = use_fused_vision_backbone
57
+ self.image_resize_strategy = image_resize_strategy
58
+
59
+ # Handle `None` default values
60
+ input_sizes = [(3, 224, 224)] if input_sizes is None else input_sizes
61
+ means = [(0.5, 0.5, 0.5)] if means is None else means
62
+ stds = [(0.5, 0.5, 0.5)] if stds is None else stds
63
+
64
+ # TIMM `data_cfg` Parameters
65
+ self.input_sizes, self.interpolations, self.means, self.stds = input_sizes, interpolations, means, stds
66
+
67
+ # Grab torchvision transforms via TIMM =>> need to parse for specific "functional" transform values!
68
+ self.tvf_resize_params, self.tvf_crop_params, self.tvf_normalize_params = [], [], []
69
+ self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None
70
+
71
+ for idx in range(len(input_sizes)):
72
+ transform = timm.data.create_transform(
73
+ input_size=self.input_sizes[idx],
74
+ interpolation=self.interpolations[idx],
75
+ mean=self.means[idx],
76
+ std=self.stds[idx],
77
+ crop_pct=1.0, # Set to 1.0 to ignore cropping (initial Resize sets `input_size`)
78
+ crop_mode="center", # Default crop mode -- no-op when `crop_pct == 1.0`
79
+ is_training=False, # No image augmentations when loading the transform!
80
+ )
81
+
82
+ # [Validation] Ensure appropriate transform structure, expected sizes
83
+ if not (
84
+ isinstance(transform, Compose)
85
+ and (len(transform.transforms) == 4)
86
+ and isinstance(transform.transforms[0], Resize)
87
+ and isinstance(transform.transforms[1], CenterCrop)
88
+ and isinstance(transform.transforms[2], ToTensor)
89
+ and isinstance(transform.transforms[3], Normalize)
90
+ and (transform.transforms[0].size == self.input_sizes[idx][-1])
91
+ and (transform.transforms[1].size == self.input_sizes[idx][-2:])
92
+ ):
93
+ raise ValueError(f"Unexpected TIMM image transformation structure/sizes: `{transform}`")
94
+
95
+ # HF Image Processors *must* be JSON-serializable; as such, cannot have torchvision. as an attribute.
96
+ # => Instead, we're going to parse the transform and call "torchvision.transforms.functional" (`tvf`)
97
+ resize_t, crop_t, norm_t = transform.transforms[0], transform.transforms[1], transform.transforms[3]
98
+ self.tvf_resize_params.append(
99
+ {
100
+ "size": resize_t.size,
101
+ "interpolation": TVF.pil_modes_mapping[resize_t.interpolation],
102
+ "max_size": None,
103
+ "antialias": True,
104
+ }
105
+ )
106
+ self.tvf_crop_params.append({"output_size": crop_t.size})
107
+ self.tvf_normalize_params.append(
108
+ {
109
+ "mean": norm_t.mean.float().numpy().tolist(),
110
+ "std": norm_t.std.float().numpy().tolist(),
111
+ "inplace": False,
112
+ }
113
+ )
114
+ self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None
115
+
116
+ # Handle Prismatic `image_resize_strategy`
117
+ if self.image_resize_strategy == "resize-naive":
118
+ self.tvf_resize_params[idx]["size"] = (resize_t.size, resize_t.size)
119
+ elif self.image_resize_strategy == "letterbox":
120
+ self.tvf_do_letterbox, self.tvf_letterbox_fill = True, tuple([int(x * 255) for x in self.means[idx]])
121
+ elif self.image_resize_strategy == "resize-crop":
122
+ pass
123
+ else:
124
+ raise ValueError(f"Image resize strategy `{self.image_resize_strategy}` is not supported!")
125
+
126
+ # Dispatch **kwargs to super()
127
+ super().__init__(**kwargs)
128
+
129
+ def apply_transform(self, img: Image.Image) -> torch.Tensor:
130
+ """Apply `functional` variant of TIMM's Transform = Compose([Resize -> CenterCrop -> ToTensor -> Normalize])"""
131
+ if self.tvf_do_letterbox:
132
+ img = letterbox_pad_transform(img, self.tvf_letterbox_fill)
133
+
134
+ # [Contract] Fused Backbones expect "channel-stacked" inputs; we'll unpack on the model side!
135
+ imgs_t = []
136
+ for idx in range(len(self.input_sizes)):
137
+ img_idx = TVF.resize(img, **self.tvf_resize_params[idx])
138
+ img_idx = TVF.center_crop(img_idx, **self.tvf_crop_params[idx])
139
+ img_idx_t = TVF.to_tensor(img_idx)
140
+ img_idx_t = TVF.normalize(img_idx_t, **self.tvf_normalize_params[idx])
141
+ imgs_t.append(img_idx_t)
142
+
143
+ # [Contract] `imgs_t` is a list of Tensors of shape [3, input_size, input_size]; stack along dim = 0
144
+ img_t = torch.vstack(imgs_t)
145
+
146
+ return img_t
147
+
148
+ def preprocess(
149
+ self,
150
+ images: Union[Image.Image, List[Image.Image]],
151
+ return_tensors: Optional[Union[str, TensorType]] = None,
152
+ **_: str,
153
+ ) -> BatchFeature:
154
+ """
155
+ Preprocess an image (or batch of images); note that unlike the `transformers :: BaseImageProcessor` we
156
+ explicitly only handle PIL.Image.Image instances for simplicity.
157
+
158
+ @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.
159
+ @param return_tensors: BatchFeature default Tensor format (e.g., "pt" for torch); if None, returns np.ndarray
160
+
161
+ @return: Instance of `transformers :: BatchFeature` with a single key "pixel_values"
162
+ """
163
+ if not isinstance(images, list):
164
+ images = [images]
165
+
166
+ # Apply `self.img_transform` to each image (will return list of torch.Tensors); stack into "batched" Tensor
167
+ pixel_values = torch.stack([self.apply_transform(img.convert("RGB")) for img in images])
168
+
169
+ # Return BatchFeature =>> note that for compatibility, constructor expects Dict[str, np.ndarray], so we convert
170
+ return BatchFeature(data={"pixel_values": pixel_values.float().numpy()}, tensor_type=return_tensors)
171
+
172
+ def __call__(self, images: Union[Image.Image, List[Image.Image]], **kwargs) -> BatchFeature:
173
+ return self.preprocess(images, **kwargs)
174
+
175
+
176
+ # === PrismaticProcessor =>> Wraps both ImageProcessor and Tokenizer ===
177
+ # =>> https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/processing_llava.py
178
+ class PrismaticProcessor(ProcessorMixin):
179
+ attributes: ClassVar[List[str]] = ["image_processor", "tokenizer"]
180
+ image_processor_class: str = "AutoImageProcessor"
181
+ tokenizer_class: str = "AutoTokenizer"
182
+
183
+ def __init__(
184
+ self,
185
+ image_processor: Optional[ImageProcessingMixin] = None,
186
+ tokenizer: Optional[PreTrainedTokenizerBase] = None,
187
+ ) -> None:
188
+ super().__init__(image_processor, tokenizer)
189
+
190
+ def __call__(
191
+ self,
192
+ text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]],
193
+ images: Union[Image.Image, List[Image.Image]],
194
+ padding: Union[bool, str, PaddingStrategy] = False,
195
+ truncation: Optional[Union[bool, str, TruncationStrategy]] = None,
196
+ max_length: Optional[int] = None,
197
+ return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
198
+ ) -> BatchFeature:
199
+ """
200
+ Preprocess a given (batch) of text/images for a Prismatic VLM; forwards text to the underlying LLM's tokenizer,
201
+ forwards images to PrismaticImageProcessor.
202
+
203
+ @param text: The (batch) of text to encode; must be a string or list of strings.
204
+ @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.
205
+ @param padding: Sequence padding strategy (if multiple specified) in < True = "longest" | "max_length" | False >
206
+ @param truncation: Truncation strategy for the output sequences; requires `max_length` to be specified
207
+ @param max_length: Maximum length (in tokens) to truncate
208
+ @param return_tensors: Type of return tensors (usually "pt" or TensorType.PYTORCH)
209
+
210
+ @return: BatchFeature with keys for `input_ids`, `attention_mask` and `pixel_values`.
211
+ """
212
+ pixel_values = self.image_processor(images, return_tensors=return_tensors)["pixel_values"]
213
+ text_inputs = self.tokenizer(
214
+ text, return_tensors=return_tensors, padding=padding, truncation=truncation, max_length=max_length
215
+ )
216
+
217
+ # [Validate] Need same number of images and text inputs!
218
+ if pixel_values.shape[0] != text_inputs.input_ids.shape[0]:
219
+ raise ValueError("Batch is malformed; expected same number of images and text inputs!")
220
+
221
+ return BatchFeature(data={**text_inputs, "pixel_values": pixel_values})
222
+
223
+ # === Tokenizer Dispatch Utilities =>> check `PreTrainedTokenizerBase` for documentation ===
224
+ def batch_decode(
225
+ self,
226
+ sequences: Union[List[int], List[List[int]], torch.Tensor, Any], # `Any` = np.ndarray | tf.Tensor
227
+ skip_special_tokens: bool = False,
228
+ clean_up_tokenization_spaces: Optional[bool] = None,
229
+ **kwargs: str,
230
+ ) -> List[str]:
231
+ return self.tokenizer.batch_decode(
232
+ sequences=sequences,
233
+ skip_special_tokens=skip_special_tokens,
234
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
235
+ **kwargs,
236
+ )
237
+
238
+ def decode(
239
+ self,
240
+ token_ids: Union[int, List[int], torch.Tensor, Any], # `Any` = np.ndarray | tf.Tensor
241
+ skip_special_tokens: bool = False,
242
+ clean_up_tokenization_spaces: Optional[bool] = None,
243
+ **kwargs: str,
244
+ ) -> str:
245
+ return self.tokenizer.decode(
246
+ token_ids=token_ids,
247
+ skip_special_tokens=skip_special_tokens,
248
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
249
+ **kwargs,
250
+ )
251
+
252
+ @property
253
+ def model_input_names(self) -> List[str]:
254
+ tokenizer_input_names = self.tokenizer.model_input_names
255
+ image_processor_input_names = self.image_processor.model_input_names
256
+
257
+ return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
processor_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "auto_map": {
3
+ "AutoProcessor": "processing_prismatic.PrismaticProcessor"
4
+ },
5
+ "processor_class": "PrismaticProcessor"
6
+ }
proprio_projector--10000_checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0c9c2da310a2284bdffe8f697ba4225128de78b7bafd41345653a478804f5036
3
+ size 67242416
special_tokens_map.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "</s>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "pad_token": {
17
+ "content": "<PAD>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "unk_token": {
24
+ "content": "<unk>",
25
+ "lstrip": false,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ }
30
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
3
+ size 499723
tokenizer_config.json ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": true,
3
+ "add_eos_token": false,
4
+ "added_tokens_decoder": {
5
+ "0": {
6
+ "content": "<unk>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "1": {
14
+ "content": "<s>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "2": {
22
+ "content": "</s>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ },
29
+ "32000": {
30
+ "content": "<PAD>",
31
+ "lstrip": false,
32
+ "normalized": false,
33
+ "rstrip": false,
34
+ "single_word": false,
35
+ "special": true
36
+ }
37
+ },
38
+ "auto_map": {
39
+ "AutoProcessor": "processing_prismatic.PrismaticProcessor"
40
+ },
41
+ "bos_token": "<s>",
42
+ "clean_up_tokenization_spaces": false,
43
+ "eos_token": "</s>",
44
+ "legacy": false,
45
+ "model_max_length": 2048,
46
+ "pad_token": "<PAD>",
47
+ "padding_side": "right",
48
+ "processor_class": "PrismaticProcessor",
49
+ "sp_model_kwargs": {},
50
+ "tokenizer_class": "LlamaTokenizer",
51
+ "unk_token": "<unk>",
52
+ "use_default_system_prompt": false
53
+ }
vision_backbone--10000_checkpoint.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:98995f617d20e54f53de85c0ec090882af4a86272be8566f1fe43afdb7c3b4b6
3
+ size 3344586069