Initial model upload
Browse files- action_head--latest_checkpoint.pt +3 -0
- added_tokens.json +4 -0
- config.json +304 -0
- config.json.back.20250403_105325 +304 -0
- config.json.back.20250403_105351 +304 -0
- config.json.back.20250403_105936 +304 -0
- config.json.back.20250403_110002 +304 -0
- config.json.back.20250403_110045 +304 -0
- config.json.back.20250403_110111 +304 -0
- config.json.back.20250403_110334 +304 -0
- config.json.back.20250403_110623 +304 -0
- config.json.back.20250403_110730 +304 -0
- config.json.back.20250403_111133 +304 -0
- config.json.back.20250403_111217 +304 -0
- config.json.back.20250403_111705 +304 -0
- config.json.back.20250403_113657 +304 -0
- config.json.back.20250403_120306 +304 -0
- config.json.back.20250403_122246 +304 -0
- config.json.back.20250403_125446 +304 -0
- config.json.back.20250403_125858 +304 -0
- config.json.back.20250403_125952 +304 -0
- config.json.back.20250403_130155 +304 -0
- config.json.back.20250403_130245 +304 -0
- config.json.back.20250403_130400 +304 -0
- config.json.back.20250403_132953 +304 -0
- config.json.back.20250403_134854 +304 -0
- configuration_prismatic.py +140 -0
- dataset_statistics.json +133 -0
- lora_adapter/README.md +202 -0
- lora_adapter/adapter_config.json +45 -0
- lora_adapter/adapter_model.safetensors +3 -0
- modeling_prismatic.py +1334 -0
- modeling_prismatic.py.back.20250403_110623 +1321 -0
- modeling_prismatic.py.back.20250403_110730 +1322 -0
- modeling_prismatic.py.back.20250403_111133 +1322 -0
- modeling_prismatic.py.back.20250403_111217 +1323 -0
- modeling_prismatic.py.back.20250403_111705 +1323 -0
- modeling_prismatic.py.back.20250403_122246 +1323 -0
- modeling_prismatic.py.back.20250403_125446 +1321 -0
- modeling_prismatic.py.back.20250403_125952 +1327 -0
- modeling_prismatic.py.back.20250403_130155 +1328 -0
- modeling_prismatic.py.back.20250403_130245 +1328 -0
- modeling_prismatic.py.back.20250403_132953 +1327 -0
- preprocessor_config.json +114 -0
- processing_prismatic.py +257 -0
- processor_config.json +6 -0
- special_tokens_map.json +39 -0
- tokenizer.json +0 -0
- tokenizer.model +3 -0
- tokenizer_config.json +64 -0
action_head--latest_checkpoint.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:fa225063fba55a7380188ad48b6cb2c0c879e5e6c953388e13522c072958298b
|
| 3 |
+
size 101219314
|
added_tokens.json
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"<ACT>": 32001,
|
| 3 |
+
"<PAD>": 32000
|
| 4 |
+
}
|
config.json
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_105325
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_105351
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_105936
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_110002
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_110045
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_110111
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_110334
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_110623
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_110730
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_111133
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_111217
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_111705
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_113657
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_120306
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_122246
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_125446
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_125858
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_125952
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_130155
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_130245
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_130400
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_132953
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
config.json.back.20250403_134854
ADDED
|
@@ -0,0 +1,304 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"norm_stats": {
|
| 3 |
+
"libero_spatial_no_noops": {
|
| 4 |
+
"action": {
|
| 5 |
+
"mean": [
|
| 6 |
+
0.15312479436397552,
|
| 7 |
+
0.13707277178764343,
|
| 8 |
+
-0.15526802837848663,
|
| 9 |
+
-0.005176450591534376,
|
| 10 |
+
-0.01120874285697937,
|
| 11 |
+
-0.020194264128804207,
|
| 12 |
+
0.4578818082809448
|
| 13 |
+
],
|
| 14 |
+
"std": [
|
| 15 |
+
0.41272708773612976,
|
| 16 |
+
0.34724321961402893,
|
| 17 |
+
0.50869220495224,
|
| 18 |
+
0.037266165018081665,
|
| 19 |
+
0.07244449853897095,
|
| 20 |
+
0.05762382969260216,
|
| 21 |
+
0.49827873706817627
|
| 22 |
+
],
|
| 23 |
+
"max": [
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.9375,
|
| 27 |
+
0.1971428543329239,
|
| 28 |
+
0.33642858266830444,
|
| 29 |
+
0.375,
|
| 30 |
+
1.0
|
| 31 |
+
],
|
| 32 |
+
"min": [
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.9375,
|
| 36 |
+
-0.1875,
|
| 37 |
+
-0.3675000071525574,
|
| 38 |
+
-0.36000001430511475,
|
| 39 |
+
0.0
|
| 40 |
+
],
|
| 41 |
+
"q01": [
|
| 42 |
+
-0.7454732114076613,
|
| 43 |
+
-0.6616071462631226,
|
| 44 |
+
-0.9375,
|
| 45 |
+
-0.1071428582072258,
|
| 46 |
+
-0.20678570866584778,
|
| 47 |
+
-0.1842857152223587,
|
| 48 |
+
0.0
|
| 49 |
+
],
|
| 50 |
+
"q99": [
|
| 51 |
+
0.9375,
|
| 52 |
+
0.8758928775787354,
|
| 53 |
+
0.9321428537368774,
|
| 54 |
+
0.1039285734295845,
|
| 55 |
+
0.17678570747375488,
|
| 56 |
+
0.14571428298950195,
|
| 57 |
+
1.0
|
| 58 |
+
],
|
| 59 |
+
"mask": [
|
| 60 |
+
true,
|
| 61 |
+
true,
|
| 62 |
+
true,
|
| 63 |
+
true,
|
| 64 |
+
true,
|
| 65 |
+
true,
|
| 66 |
+
false
|
| 67 |
+
]
|
| 68 |
+
},
|
| 69 |
+
"proprio": {
|
| 70 |
+
"mean": [
|
| 71 |
+
-0.024462558329105377,
|
| 72 |
+
0.106529600918293,
|
| 73 |
+
1.0580483675003052,
|
| 74 |
+
3.0628468990325928,
|
| 75 |
+
-0.10464039444923401,
|
| 76 |
+
0.08307311683893204,
|
| 77 |
+
0.01995457336306572,
|
| 78 |
+
-0.020162804052233696
|
| 79 |
+
],
|
| 80 |
+
"std": [
|
| 81 |
+
0.1101478561758995,
|
| 82 |
+
0.13784688711166382,
|
| 83 |
+
0.1044282391667366,
|
| 84 |
+
0.10451053828001022,
|
| 85 |
+
0.4112098217010498,
|
| 86 |
+
0.2176690548658371,
|
| 87 |
+
0.017260896041989326,
|
| 88 |
+
0.0171116404235363
|
| 89 |
+
],
|
| 90 |
+
"max": [
|
| 91 |
+
0.1759040206670761,
|
| 92 |
+
0.3904820382595062,
|
| 93 |
+
1.3290715217590332,
|
| 94 |
+
3.4566118717193604,
|
| 95 |
+
1.2268599271774292,
|
| 96 |
+
1.0429412126541138,
|
| 97 |
+
0.041053611785173416,
|
| 98 |
+
0.000775813648942858
|
| 99 |
+
],
|
| 100 |
+
"min": [
|
| 101 |
+
-0.3095473051071167,
|
| 102 |
+
-0.29250794649124146,
|
| 103 |
+
0.9095591306686401,
|
| 104 |
+
2.497488260269165,
|
| 105 |
+
-1.8006486892700195,
|
| 106 |
+
-0.7207611203193665,
|
| 107 |
+
-0.0004703797458205372,
|
| 108 |
+
-0.041536275297403336
|
| 109 |
+
],
|
| 110 |
+
"q01": [
|
| 111 |
+
-0.2727657300233841,
|
| 112 |
+
-0.23721413239836692,
|
| 113 |
+
0.9160063165426254,
|
| 114 |
+
2.77949666261673,
|
| 115 |
+
-1.3187511622905732,
|
| 116 |
+
-0.41989982962608335,
|
| 117 |
+
0.001503719249740243,
|
| 118 |
+
-0.03989770736545324
|
| 119 |
+
],
|
| 120 |
+
"q99": [
|
| 121 |
+
0.13529365032911292,
|
| 122 |
+
0.3629165390133857,
|
| 123 |
+
1.2862326657772063,
|
| 124 |
+
3.2829698753356933,
|
| 125 |
+
0.9332760351896285,
|
| 126 |
+
0.6325724506378171,
|
| 127 |
+
0.039933966137468815,
|
| 128 |
+
-0.001671919699292631
|
| 129 |
+
]
|
| 130 |
+
},
|
| 131 |
+
"num_transitions": 52970,
|
| 132 |
+
"num_trajectories": 432
|
| 133 |
+
}
|
| 134 |
+
},
|
| 135 |
+
"n_action_bins": 256,
|
| 136 |
+
"vision_backbone_id": "dinosiglip-vit-so-224px",
|
| 137 |
+
"llm_backbone_id": "llama2-7b-pure",
|
| 138 |
+
"arch_specifier": "no-align+fused-gelu-mlp",
|
| 139 |
+
"output_projector_states": false,
|
| 140 |
+
"use_fused_vision_backbone": true,
|
| 141 |
+
"timm_model_ids": [
|
| 142 |
+
"vit_large_patch14_reg4_dinov2.lvd142m",
|
| 143 |
+
"vit_so400m_patch14_siglip_224"
|
| 144 |
+
],
|
| 145 |
+
"timm_override_act_layers": [
|
| 146 |
+
null,
|
| 147 |
+
null
|
| 148 |
+
],
|
| 149 |
+
"image_sizes": [
|
| 150 |
+
224,
|
| 151 |
+
224
|
| 152 |
+
],
|
| 153 |
+
"image_resize_strategy": "resize-naive",
|
| 154 |
+
"hf_llm_id": "meta-llama/Llama-2-7b-hf",
|
| 155 |
+
"llm_max_length": 2048,
|
| 156 |
+
"pad_token_id": 32000,
|
| 157 |
+
"pad_to_multiple_of": 64,
|
| 158 |
+
"text_config": {
|
| 159 |
+
"vocab_size": 32064,
|
| 160 |
+
"max_position_embeddings": 2048,
|
| 161 |
+
"hidden_size": 4096,
|
| 162 |
+
"intermediate_size": 11008,
|
| 163 |
+
"num_hidden_layers": 32,
|
| 164 |
+
"num_attention_heads": 32,
|
| 165 |
+
"num_key_value_heads": 32,
|
| 166 |
+
"hidden_act": "silu",
|
| 167 |
+
"initializer_range": 0.02,
|
| 168 |
+
"rms_norm_eps": 1e-06,
|
| 169 |
+
"pretraining_tp": 1,
|
| 170 |
+
"use_cache": true,
|
| 171 |
+
"rope_theta": 10000.0,
|
| 172 |
+
"rope_scaling": null,
|
| 173 |
+
"attention_bias": false,
|
| 174 |
+
"attention_dropout": 0.0,
|
| 175 |
+
"return_dict": true,
|
| 176 |
+
"output_hidden_states": false,
|
| 177 |
+
"output_attentions": false,
|
| 178 |
+
"torchscript": false,
|
| 179 |
+
"torch_dtype": "bfloat16",
|
| 180 |
+
"use_bfloat16": false,
|
| 181 |
+
"tf_legacy_loss": false,
|
| 182 |
+
"pruned_heads": {},
|
| 183 |
+
"tie_word_embeddings": false,
|
| 184 |
+
"chunk_size_feed_forward": 0,
|
| 185 |
+
"is_encoder_decoder": false,
|
| 186 |
+
"is_decoder": false,
|
| 187 |
+
"cross_attention_hidden_size": null,
|
| 188 |
+
"add_cross_attention": false,
|
| 189 |
+
"tie_encoder_decoder": false,
|
| 190 |
+
"max_length": 20,
|
| 191 |
+
"min_length": 0,
|
| 192 |
+
"do_sample": false,
|
| 193 |
+
"early_stopping": false,
|
| 194 |
+
"num_beams": 1,
|
| 195 |
+
"num_beam_groups": 1,
|
| 196 |
+
"diversity_penalty": 0.0,
|
| 197 |
+
"temperature": 1.0,
|
| 198 |
+
"top_k": 50,
|
| 199 |
+
"top_p": 1.0,
|
| 200 |
+
"typical_p": 1.0,
|
| 201 |
+
"repetition_penalty": 1.0,
|
| 202 |
+
"length_penalty": 1.0,
|
| 203 |
+
"no_repeat_ngram_size": 0,
|
| 204 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 205 |
+
"bad_words_ids": null,
|
| 206 |
+
"num_return_sequences": 1,
|
| 207 |
+
"output_scores": false,
|
| 208 |
+
"return_dict_in_generate": false,
|
| 209 |
+
"forced_bos_token_id": null,
|
| 210 |
+
"forced_eos_token_id": null,
|
| 211 |
+
"remove_invalid_values": false,
|
| 212 |
+
"exponential_decay_length_penalty": null,
|
| 213 |
+
"suppress_tokens": null,
|
| 214 |
+
"begin_suppress_tokens": null,
|
| 215 |
+
"architectures": null,
|
| 216 |
+
"finetuning_task": null,
|
| 217 |
+
"id2label": {
|
| 218 |
+
"0": "LABEL_0",
|
| 219 |
+
"1": "LABEL_1"
|
| 220 |
+
},
|
| 221 |
+
"label2id": {
|
| 222 |
+
"LABEL_0": 0,
|
| 223 |
+
"LABEL_1": 1
|
| 224 |
+
},
|
| 225 |
+
"tokenizer_class": null,
|
| 226 |
+
"prefix": null,
|
| 227 |
+
"bos_token_id": 1,
|
| 228 |
+
"pad_token_id": 32000,
|
| 229 |
+
"eos_token_id": 2,
|
| 230 |
+
"sep_token_id": null,
|
| 231 |
+
"decoder_start_token_id": null,
|
| 232 |
+
"task_specific_params": null,
|
| 233 |
+
"problem_type": null,
|
| 234 |
+
"_name_or_path": "",
|
| 235 |
+
"model_type": "llama"
|
| 236 |
+
},
|
| 237 |
+
"return_dict": true,
|
| 238 |
+
"output_hidden_states": false,
|
| 239 |
+
"output_attentions": false,
|
| 240 |
+
"torchscript": false,
|
| 241 |
+
"torch_dtype": "bfloat16",
|
| 242 |
+
"use_bfloat16": false,
|
| 243 |
+
"tf_legacy_loss": false,
|
| 244 |
+
"pruned_heads": {},
|
| 245 |
+
"tie_word_embeddings": true,
|
| 246 |
+
"chunk_size_feed_forward": 0,
|
| 247 |
+
"is_encoder_decoder": false,
|
| 248 |
+
"is_decoder": false,
|
| 249 |
+
"cross_attention_hidden_size": null,
|
| 250 |
+
"add_cross_attention": false,
|
| 251 |
+
"tie_encoder_decoder": false,
|
| 252 |
+
"max_length": 20,
|
| 253 |
+
"min_length": 0,
|
| 254 |
+
"do_sample": false,
|
| 255 |
+
"early_stopping": false,
|
| 256 |
+
"num_beams": 1,
|
| 257 |
+
"num_beam_groups": 1,
|
| 258 |
+
"diversity_penalty": 0.0,
|
| 259 |
+
"temperature": 1.0,
|
| 260 |
+
"top_k": 50,
|
| 261 |
+
"top_p": 1.0,
|
| 262 |
+
"typical_p": 1.0,
|
| 263 |
+
"repetition_penalty": 1.0,
|
| 264 |
+
"length_penalty": 1.0,
|
| 265 |
+
"no_repeat_ngram_size": 0,
|
| 266 |
+
"encoder_no_repeat_ngram_size": 0,
|
| 267 |
+
"bad_words_ids": null,
|
| 268 |
+
"num_return_sequences": 1,
|
| 269 |
+
"output_scores": false,
|
| 270 |
+
"return_dict_in_generate": false,
|
| 271 |
+
"forced_bos_token_id": null,
|
| 272 |
+
"forced_eos_token_id": null,
|
| 273 |
+
"remove_invalid_values": false,
|
| 274 |
+
"exponential_decay_length_penalty": null,
|
| 275 |
+
"suppress_tokens": null,
|
| 276 |
+
"begin_suppress_tokens": null,
|
| 277 |
+
"architectures": [
|
| 278 |
+
"OpenVLAForActionPrediction"
|
| 279 |
+
],
|
| 280 |
+
"finetuning_task": null,
|
| 281 |
+
"id2label": {
|
| 282 |
+
"0": "LABEL_0",
|
| 283 |
+
"1": "LABEL_1"
|
| 284 |
+
},
|
| 285 |
+
"label2id": {
|
| 286 |
+
"LABEL_0": 0,
|
| 287 |
+
"LABEL_1": 1
|
| 288 |
+
},
|
| 289 |
+
"tokenizer_class": null,
|
| 290 |
+
"prefix": null,
|
| 291 |
+
"bos_token_id": null,
|
| 292 |
+
"eos_token_id": null,
|
| 293 |
+
"sep_token_id": null,
|
| 294 |
+
"decoder_start_token_id": null,
|
| 295 |
+
"task_specific_params": null,
|
| 296 |
+
"problem_type": null,
|
| 297 |
+
"_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 298 |
+
"transformers_version": "4.40.1",
|
| 299 |
+
"auto_map": {
|
| 300 |
+
"AutoConfig": "configuration_prismatic.OpenVLAConfig",
|
| 301 |
+
"AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction"
|
| 302 |
+
},
|
| 303 |
+
"model_type": "openvla"
|
| 304 |
+
}
|
configuration_prismatic.py
ADDED
|
@@ -0,0 +1,140 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
VISION_BACKBONE_TO_RESOLUTION: Dict[str, List[int]] = {
|
| 16 |
+
"clip-vit-l": [224], "siglip-vit-so400m": [224], "dinov2-vit-l": [224], "in1k-vit-l": [224],
|
| 17 |
+
|
| 18 |
+
"clip-vit-l-336px": [336],
|
| 19 |
+
"siglip-vit-so400m-384px": [384],
|
| 20 |
+
|
| 21 |
+
"dinoclip-vit-l-336px": [336, 336],
|
| 22 |
+
"dinosiglip-vit-so-224px": [224, 224],
|
| 23 |
+
"dinosiglip-vit-so-384px": [384, 384],
|
| 24 |
+
}
|
| 25 |
+
VISION_BACKBONE_TO_TIMM_ID: Dict[str, List[str]] = {
|
| 26 |
+
"clip-vit-l": ["vit_large_patch14_clip_224.openai"],
|
| 27 |
+
"clip-vit-l-336px": ["vit_large_patch14_clip_336.openai"],
|
| 28 |
+
|
| 29 |
+
"dinov2-vit-l": ["vit_large_patch14_reg4_dinov2.lvd142m"],
|
| 30 |
+
"in1k-vit-l": ["vit_large_patch16_224.augreg_in21k_ft_in1k"],
|
| 31 |
+
|
| 32 |
+
"siglip-vit-so400m": ["vit_so400m_patch14_siglip_224"],
|
| 33 |
+
"siglip-vit-so400m-384px": ["vit_so400m_patch14_siglip_384"],
|
| 34 |
+
|
| 35 |
+
"dinoclip-vit-l-336px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_large_patch14_clip_336.openai"],
|
| 36 |
+
"dinosiglip-vit-so-224px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_224"],
|
| 37 |
+
"dinosiglip-vit-so-384px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_384"],
|
| 38 |
+
}
|
| 39 |
+
TIMM_OVERRIDE_ACT_LAYER: Dict[str, List[Optional[str]]] = {
|
| 40 |
+
"clip-vit-l": ["quick_gelu"], "clip-vit-l-336px": ["quick_gelu"],
|
| 41 |
+
"dinov2-vit-l": [None], "in1k-vit-l": [None],
|
| 42 |
+
"siglip-vit-so400m": [None], "siglip-vit-so400m-384px": [None],
|
| 43 |
+
"dinoclip-vit-l-336px": [None, "quick_gelu"],
|
| 44 |
+
"dinosiglip-vit-so-224px": [None, None], "dinosiglip-vit-so-384px": [None, None]
|
| 45 |
+
}
|
| 46 |
+
|
| 47 |
+
LLM_BACKBONE_TO_HF_PATH = {
|
| 48 |
+
"llama2-7b-pure": "meta-llama/Llama-2-7b-hf", "llama2-13b-pure": "meta-llama/Llama-2-13b-hf",
|
| 49 |
+
"llama2-7b-chat": "meta-llama/Llama-2-7b-chat-hf", "llama2-13b-chat": "meta-llama/Llama-2-13b-chat-hf",
|
| 50 |
+
|
| 51 |
+
"vicuna-v15-7b": "lmsys/vicuna-7b-v1.5", "vicuna-v15-13b": "lmsys/vicuna-13b-v1.5",
|
| 52 |
+
|
| 53 |
+
"mistral-v0.1-7b-pure": "mistralai/Mistral-7B-v0.1",
|
| 54 |
+
"mistral-v0.1-7b-instruct": "mistralai/Mistral-7B-Instruct-v0.1",
|
| 55 |
+
|
| 56 |
+
"phi-2-3b": "microsoft/phi-2",
|
| 57 |
+
}
|
| 58 |
+
LLM_BACKBONE_TO_HF_METACLASS = {
|
| 59 |
+
"llama2-7b-pure": "llama", "llama2-13b-pure": "llama", "llama2-7b-chat": "llama", "llama2-13b-chat": "llama",
|
| 60 |
+
"vicuna-v15-7b": "llama", "vicuna-v15-13b": "llama",
|
| 61 |
+
|
| 62 |
+
"mistral-v0.1-7b-pure": "mistral", "mistral-v0.1-7b-instruct": "mistral",
|
| 63 |
+
|
| 64 |
+
"phi-2-3b": "phi",
|
| 65 |
+
}
|
| 66 |
+
|
| 67 |
+
VALID_VISION_BACKBONES = set(VISION_BACKBONE_TO_RESOLUTION.keys())
|
| 68 |
+
VALID_LLM_BACKBONES = set(LLM_BACKBONE_TO_HF_PATH)
|
| 69 |
+
# fmt: on
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
class PrismaticConfig(PretrainedConfig):
|
| 73 |
+
model_type: str = "prismatic"
|
| 74 |
+
is_composition: bool = False
|
| 75 |
+
|
| 76 |
+
def __init__(
|
| 77 |
+
self,
|
| 78 |
+
vision_backbone_id: str = "siglip-vit-so400m",
|
| 79 |
+
llm_backbone_id: str = "vicuna-v15-7b",
|
| 80 |
+
arch_specifier: str = "no-align+gelu-mlp",
|
| 81 |
+
use_fused_vision_backbone: Optional[bool] = None,
|
| 82 |
+
image_resize_strategy: str = "letterbox",
|
| 83 |
+
text_config: Optional[Dict[str, Any]] = None,
|
| 84 |
+
llm_max_length: int = 2048,
|
| 85 |
+
pad_token_id: int = 32000,
|
| 86 |
+
pad_to_multiple_of: int = 64,
|
| 87 |
+
output_projector_states: bool = False,
|
| 88 |
+
**kwargs: str,
|
| 89 |
+
) -> None:
|
| 90 |
+
if vision_backbone_id not in VALID_VISION_BACKBONES:
|
| 91 |
+
raise ValueError(f"Vision backbone `{vision_backbone_id}` not in {VALID_VISION_BACKBONES = }")
|
| 92 |
+
|
| 93 |
+
if llm_backbone_id not in VALID_LLM_BACKBONES:
|
| 94 |
+
raise ValueError(f"LLM backbone `{llm_backbone_id}` not in {VALID_LLM_BACKBONES = }")
|
| 95 |
+
|
| 96 |
+
# Set Prismatic Configuration Fields
|
| 97 |
+
self.vision_backbone_id = vision_backbone_id
|
| 98 |
+
self.llm_backbone_id = llm_backbone_id
|
| 99 |
+
self.arch_specifier = arch_specifier
|
| 100 |
+
self.output_projector_states = output_projector_states
|
| 101 |
+
|
| 102 |
+
# [Contract] All vision backbone parameters are lists =>> supports fused backbones with different preprocessing
|
| 103 |
+
self.use_fused_vision_backbone = (
|
| 104 |
+
use_fused_vision_backbone
|
| 105 |
+
if use_fused_vision_backbone is not None
|
| 106 |
+
else any(self.vision_backbone_id.startswith(v) for v in ["dinoclip", "dinosiglip"])
|
| 107 |
+
)
|
| 108 |
+
|
| 109 |
+
self.timm_model_ids = VISION_BACKBONE_TO_TIMM_ID[self.vision_backbone_id]
|
| 110 |
+
self.timm_override_act_layers = TIMM_OVERRIDE_ACT_LAYER[self.vision_backbone_id]
|
| 111 |
+
self.image_sizes = VISION_BACKBONE_TO_RESOLUTION[self.vision_backbone_id]
|
| 112 |
+
self.image_resize_strategy = image_resize_strategy
|
| 113 |
+
|
| 114 |
+
self.hf_llm_id = LLM_BACKBONE_TO_HF_PATH[self.llm_backbone_id]
|
| 115 |
+
self.llm_max_length = llm_max_length
|
| 116 |
+
self.pad_token_id, self.pad_to_multiple_of = pad_token_id, pad_to_multiple_of
|
| 117 |
+
|
| 118 |
+
# [IMPORTANT] HF Utilities actually look for a `text_config` field... we need to use that specific naming!
|
| 119 |
+
self.text_config = (
|
| 120 |
+
CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]](**text_config)
|
| 121 |
+
if text_config is not None
|
| 122 |
+
else CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]]()
|
| 123 |
+
)
|
| 124 |
+
|
| 125 |
+
# Dispatch **kwargs to super() =>> note that `pad_token_id` collides, so we pass it in here as well...
|
| 126 |
+
super().__init__(pad_token_id=pad_token_id, **kwargs)
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
class OpenVLAConfig(PrismaticConfig):
|
| 130 |
+
model_type: str = "openvla"
|
| 131 |
+
|
| 132 |
+
def __init__(
|
| 133 |
+
self,
|
| 134 |
+
norm_stats: Optional[Dict[str, Dict[str, Dict[str, Dict[str, List[float]]]]]] = None,
|
| 135 |
+
n_action_bins: int = 256,
|
| 136 |
+
**kwargs: str,
|
| 137 |
+
) -> None:
|
| 138 |
+
self.norm_stats, self.n_action_bins = norm_stats, n_action_bins
|
| 139 |
+
|
| 140 |
+
super().__init__(**kwargs)
|
dataset_statistics.json
ADDED
|
@@ -0,0 +1,133 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"libero_spatial_no_noops": {
|
| 3 |
+
"action": {
|
| 4 |
+
"mean": [
|
| 5 |
+
0.15312479436397552,
|
| 6 |
+
0.13707277178764343,
|
| 7 |
+
-0.15526802837848663,
|
| 8 |
+
-0.005176450591534376,
|
| 9 |
+
-0.01120874285697937,
|
| 10 |
+
-0.020194264128804207,
|
| 11 |
+
0.4578818082809448
|
| 12 |
+
],
|
| 13 |
+
"std": [
|
| 14 |
+
0.41272708773612976,
|
| 15 |
+
0.34724321961402893,
|
| 16 |
+
0.50869220495224,
|
| 17 |
+
0.037266165018081665,
|
| 18 |
+
0.07244449853897095,
|
| 19 |
+
0.05762382969260216,
|
| 20 |
+
0.49827873706817627
|
| 21 |
+
],
|
| 22 |
+
"max": [
|
| 23 |
+
0.9375,
|
| 24 |
+
0.9375,
|
| 25 |
+
0.9375,
|
| 26 |
+
0.1971428543329239,
|
| 27 |
+
0.33642858266830444,
|
| 28 |
+
0.375,
|
| 29 |
+
1.0
|
| 30 |
+
],
|
| 31 |
+
"min": [
|
| 32 |
+
-0.9375,
|
| 33 |
+
-0.9375,
|
| 34 |
+
-0.9375,
|
| 35 |
+
-0.1875,
|
| 36 |
+
-0.3675000071525574,
|
| 37 |
+
-0.36000001430511475,
|
| 38 |
+
0.0
|
| 39 |
+
],
|
| 40 |
+
"q01": [
|
| 41 |
+
-0.7454732114076613,
|
| 42 |
+
-0.6616071462631226,
|
| 43 |
+
-0.9375,
|
| 44 |
+
-0.1071428582072258,
|
| 45 |
+
-0.20678570866584778,
|
| 46 |
+
-0.1842857152223587,
|
| 47 |
+
0.0
|
| 48 |
+
],
|
| 49 |
+
"q99": [
|
| 50 |
+
0.9375,
|
| 51 |
+
0.8758928775787354,
|
| 52 |
+
0.9321428537368774,
|
| 53 |
+
0.1039285734295845,
|
| 54 |
+
0.17678570747375488,
|
| 55 |
+
0.14571428298950195,
|
| 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.024462558329105377,
|
| 71 |
+
0.106529600918293,
|
| 72 |
+
1.0580483675003052,
|
| 73 |
+
3.0628468990325928,
|
| 74 |
+
-0.10464039444923401,
|
| 75 |
+
0.08307311683893204,
|
| 76 |
+
0.01995457336306572,
|
| 77 |
+
-0.020162804052233696
|
| 78 |
+
],
|
| 79 |
+
"std": [
|
| 80 |
+
0.1101478561758995,
|
| 81 |
+
0.13784688711166382,
|
| 82 |
+
0.1044282391667366,
|
| 83 |
+
0.10451053828001022,
|
| 84 |
+
0.4112098217010498,
|
| 85 |
+
0.2176690548658371,
|
| 86 |
+
0.017260896041989326,
|
| 87 |
+
0.0171116404235363
|
| 88 |
+
],
|
| 89 |
+
"max": [
|
| 90 |
+
0.1759040206670761,
|
| 91 |
+
0.3904820382595062,
|
| 92 |
+
1.3290715217590332,
|
| 93 |
+
3.4566118717193604,
|
| 94 |
+
1.2268599271774292,
|
| 95 |
+
1.0429412126541138,
|
| 96 |
+
0.041053611785173416,
|
| 97 |
+
0.000775813648942858
|
| 98 |
+
],
|
| 99 |
+
"min": [
|
| 100 |
+
-0.3095473051071167,
|
| 101 |
+
-0.29250794649124146,
|
| 102 |
+
0.9095591306686401,
|
| 103 |
+
2.497488260269165,
|
| 104 |
+
-1.8006486892700195,
|
| 105 |
+
-0.7207611203193665,
|
| 106 |
+
-0.0004703797458205372,
|
| 107 |
+
-0.041536275297403336
|
| 108 |
+
],
|
| 109 |
+
"q01": [
|
| 110 |
+
-0.2727657300233841,
|
| 111 |
+
-0.23721413239836692,
|
| 112 |
+
0.9160063165426254,
|
| 113 |
+
2.77949666261673,
|
| 114 |
+
-1.3187511622905732,
|
| 115 |
+
-0.41989982962608335,
|
| 116 |
+
0.001503719249740243,
|
| 117 |
+
-0.03989770736545324
|
| 118 |
+
],
|
| 119 |
+
"q99": [
|
| 120 |
+
0.13529365032911292,
|
| 121 |
+
0.3629165390133857,
|
| 122 |
+
1.2862326657772063,
|
| 123 |
+
3.2829698753356933,
|
| 124 |
+
0.9332760351896285,
|
| 125 |
+
0.6325724506378171,
|
| 126 |
+
0.039933966137468815,
|
| 127 |
+
-0.001671919699292631
|
| 128 |
+
]
|
| 129 |
+
},
|
| 130 |
+
"num_transitions": 52970,
|
| 131 |
+
"num_trajectories": 432
|
| 132 |
+
}
|
| 133 |
+
}
|
lora_adapter/README.md
ADDED
|
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
base_model: /home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0
|
| 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,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"alpha_pattern": {},
|
| 3 |
+
"auto_mapping": {
|
| 4 |
+
"base_model_class": "OpenVLAForActionPrediction",
|
| 5 |
+
"parent_library": "transformers_modules.31f090d05236101ebfc381b61c674dd4746d4ce0.modeling_prismatic"
|
| 6 |
+
},
|
| 7 |
+
"base_model_name_or_path": "/home/user1/.cache/huggingface/hub/models--openvla--openvla-7b/snapshots/31f090d05236101ebfc381b61c674dd4746d4ce0",
|
| 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": null,
|
| 21 |
+
"peft_type": "LORA",
|
| 22 |
+
"r": 32,
|
| 23 |
+
"rank_pattern": {},
|
| 24 |
+
"revision": null,
|
| 25 |
+
"target_modules": [
|
| 26 |
+
"v_proj",
|
| 27 |
+
"q",
|
| 28 |
+
"k_proj",
|
| 29 |
+
"qkv",
|
| 30 |
+
"fc2",
|
| 31 |
+
"fc3",
|
| 32 |
+
"down_proj",
|
| 33 |
+
"up_proj",
|
| 34 |
+
"fc1",
|
| 35 |
+
"o_proj",
|
| 36 |
+
"lm_head",
|
| 37 |
+
"proj",
|
| 38 |
+
"q_proj",
|
| 39 |
+
"kv",
|
| 40 |
+
"gate_proj"
|
| 41 |
+
],
|
| 42 |
+
"task_type": null,
|
| 43 |
+
"use_dora": false,
|
| 44 |
+
"use_rslora": false
|
| 45 |
+
}
|
lora_adapter/adapter_model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:4d35099b6525c776502a29ac76c738068b21d5d8808c9a4249ff178be02ff75e
|
| 3 |
+
size 484458600
|
modeling_prismatic.py
ADDED
|
@@ -0,0 +1,1334 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 950 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 951 |
+
)
|
| 952 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 953 |
+
# first language_model_output.hidden_states , is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 954 |
+
if self.preset:
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
actions_hidden_states = language_model_output.hidden_states[0][-1]
|
| 957 |
+
actions_hidden_states = actions_hidden_states[:, -1]
|
| 958 |
+
else:
|
| 959 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2,output_hidden_states=True,return_dict_in_generate=True)
|
| 960 |
+
actions_hidden_states = language_model_output.hidden_states[1][-1]
|
| 961 |
+
actions_hidden_states = actions_hidden_states[:, -1]
|
| 962 |
+
|
| 963 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 964 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 965 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 966 |
+
|
| 967 |
+
return normalized_actions, actions_hidden_states
|
| 968 |
+
|
| 969 |
+
def predict_action(
|
| 970 |
+
self,
|
| 971 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 972 |
+
unnorm_key: Optional[str] = None,
|
| 973 |
+
proprio=None,
|
| 974 |
+
proprio_projector=None,
|
| 975 |
+
action_head=None,
|
| 976 |
+
noisy_action_projector=None,
|
| 977 |
+
use_film: bool = False,
|
| 978 |
+
**kwargs: str,
|
| 979 |
+
) -> np.ndarray:
|
| 980 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 981 |
+
|
| 982 |
+
Args:
|
| 983 |
+
input_ids: Input token ids
|
| 984 |
+
unnorm_key: Key for unnormalization statistics
|
| 985 |
+
proprio: Proprioceptive features
|
| 986 |
+
proprio_projector: Projector for proprioceptive features
|
| 987 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 988 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 989 |
+
use_film: Whether to use FiLM conditioning
|
| 990 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 991 |
+
|
| 992 |
+
Returns:
|
| 993 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 994 |
+
"""
|
| 995 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 996 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 997 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 998 |
+
input_ids = torch.cat(
|
| 999 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 1000 |
+
)
|
| 1001 |
+
|
| 1002 |
+
pixel_values = kwargs["pixel_values"]
|
| 1003 |
+
attention_mask = kwargs["attention_mask"]
|
| 1004 |
+
|
| 1005 |
+
# Create fake labels tensor (needed for action mask)
|
| 1006 |
+
labels = input_ids.clone()
|
| 1007 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1008 |
+
|
| 1009 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1010 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1011 |
+
|
| 1012 |
+
# Prepare inputs by adding necessary tokens
|
| 1013 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1014 |
+
|
| 1015 |
+
# Update labels tensor for action mask computation later
|
| 1016 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1017 |
+
|
| 1018 |
+
# Get input embeddings and action masks
|
| 1019 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1020 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1021 |
+
|
| 1022 |
+
# Extract language embeddings
|
| 1023 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1024 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1025 |
+
)
|
| 1026 |
+
|
| 1027 |
+
# Process vision features
|
| 1028 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1029 |
+
|
| 1030 |
+
# Add proprioceptive features if provided
|
| 1031 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1032 |
+
if use_proprio:
|
| 1033 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1034 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1035 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1036 |
+
)
|
| 1037 |
+
|
| 1038 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1039 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1040 |
+
|
| 1041 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1042 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1043 |
+
if use_proprio:
|
| 1044 |
+
NUM_PATCHES += 1
|
| 1045 |
+
|
| 1046 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1047 |
+
input_embeddings,
|
| 1048 |
+
all_actions_mask,
|
| 1049 |
+
projected_patch_embeddings,
|
| 1050 |
+
attention_mask,
|
| 1051 |
+
labels,
|
| 1052 |
+
NUM_PATCHES,
|
| 1053 |
+
NUM_PROMPT_TOKENS,
|
| 1054 |
+
action_head,
|
| 1055 |
+
)
|
| 1056 |
+
|
| 1057 |
+
# Unnormalize predicted actions
|
| 1058 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1059 |
+
|
| 1060 |
+
return actions, actions_hidden_states
|
| 1061 |
+
|
| 1062 |
+
@staticmethod
|
| 1063 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1064 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1065 |
+
if unnorm_key is None:
|
| 1066 |
+
assert len(norm_stats) == 1, (
|
| 1067 |
+
f"Your model was trained on more than one dataset, "
|
| 1068 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1069 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1070 |
+
)
|
| 1071 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1072 |
+
# norm states没有加载libero, 为什么?
|
| 1073 |
+
assert unnorm_key in norm_stats, (
|
| 1074 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1075 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1076 |
+
)
|
| 1077 |
+
return unnorm_key
|
| 1078 |
+
|
| 1079 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1080 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1081 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1082 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1083 |
+
|
| 1084 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1085 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1086 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1087 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1088 |
+
|
| 1089 |
+
|
| 1090 |
+
def lisa_forward(
|
| 1091 |
+
self,
|
| 1092 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1093 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1094 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1095 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1096 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1097 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1098 |
+
use_cache: Optional[bool] = None,
|
| 1099 |
+
output_attentions: Optional[bool] = None,
|
| 1100 |
+
output_hidden_states: Optional[bool] = None,
|
| 1101 |
+
output_projector_features: Optional[bool] = None,
|
| 1102 |
+
return_dict: Optional[bool] = None,
|
| 1103 |
+
proprio=None,
|
| 1104 |
+
proprio_projector=None,
|
| 1105 |
+
**kwargs
|
| 1106 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1107 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1108 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1109 |
+
output_hidden_states = (
|
| 1110 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1111 |
+
)
|
| 1112 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1113 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1114 |
+
|
| 1115 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1116 |
+
use_cache = use_cache and not self.training
|
| 1117 |
+
|
| 1118 |
+
# Instantiate Placeholder for Projector Features
|
| 1119 |
+
projected_patch_embeddings = None
|
| 1120 |
+
|
| 1121 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1122 |
+
if input_ids.shape[1] == 1:
|
| 1123 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1124 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1125 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1126 |
+
|
| 1127 |
+
language_model_output = self.language_model(
|
| 1128 |
+
input_ids=input_ids,
|
| 1129 |
+
attention_mask=None,
|
| 1130 |
+
position_ids=None,
|
| 1131 |
+
past_key_values=past_key_values,
|
| 1132 |
+
inputs_embeds=None,
|
| 1133 |
+
labels=None,
|
| 1134 |
+
use_cache=use_cache,
|
| 1135 |
+
output_attentions=output_attentions,
|
| 1136 |
+
output_hidden_states=output_hidden_states,
|
| 1137 |
+
return_dict=return_dict,
|
| 1138 |
+
)
|
| 1139 |
+
|
| 1140 |
+
# === Handle Unimodal Forward ===
|
| 1141 |
+
elif pixel_values is None:
|
| 1142 |
+
raise NotImplementedError
|
| 1143 |
+
|
| 1144 |
+
# === Handle Multimodal Forward ===
|
| 1145 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1146 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1147 |
+
|
| 1148 |
+
# Get input embeddings (from language model embeddings)
|
| 1149 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1150 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1151 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1152 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1153 |
+
# ) # (B, lang_seq_len, llm_dim) 这里���会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1154 |
+
# Get visual features
|
| 1155 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1156 |
+
|
| 1157 |
+
# Add proprioceptive state if provided
|
| 1158 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1159 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1160 |
+
)
|
| 1161 |
+
|
| 1162 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1163 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1164 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1165 |
+
|
| 1166 |
+
# Build multimodal embeddings & attention mask
|
| 1167 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1168 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1169 |
+
)
|
| 1170 |
+
|
| 1171 |
+
# Build labels for multimodal sequence if needed
|
| 1172 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1173 |
+
|
| 1174 |
+
# Dispatch to language model
|
| 1175 |
+
language_model_output = self.language_model(
|
| 1176 |
+
input_ids=None,
|
| 1177 |
+
attention_mask=multimodal_attention_mask,
|
| 1178 |
+
position_ids=None,
|
| 1179 |
+
past_key_values=None,
|
| 1180 |
+
inputs_embeds=multimodal_embeddings,
|
| 1181 |
+
labels=multimodal_labels,
|
| 1182 |
+
use_cache=use_cache,
|
| 1183 |
+
output_attentions=output_attentions,
|
| 1184 |
+
output_hidden_states=output_hidden_states,
|
| 1185 |
+
return_dict=return_dict,
|
| 1186 |
+
)
|
| 1187 |
+
|
| 1188 |
+
|
| 1189 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1190 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1191 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1192 |
+
|
| 1193 |
+
else:
|
| 1194 |
+
raise ValueError(
|
| 1195 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1196 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1197 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1198 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1199 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1200 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1201 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1202 |
+
f"=> `use_cache` = {use_cache}"
|
| 1203 |
+
)
|
| 1204 |
+
|
| 1205 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1206 |
+
if not return_dict:
|
| 1207 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1208 |
+
return *language_model_output, projected_patch_embeddings
|
| 1209 |
+
|
| 1210 |
+
return language_model_output
|
| 1211 |
+
|
| 1212 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1213 |
+
loss=language_model_output.loss,
|
| 1214 |
+
logits=language_model_output.logits,
|
| 1215 |
+
past_key_values=language_model_output.past_key_values,
|
| 1216 |
+
hidden_states=language_model_output.hidden_states,
|
| 1217 |
+
attentions=language_model_output.attentions,
|
| 1218 |
+
projector_features=projected_patch_embeddings,
|
| 1219 |
+
)
|
| 1220 |
+
|
| 1221 |
+
def lisa_predict_action(
|
| 1222 |
+
self,
|
| 1223 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1224 |
+
unnorm_key: Optional[str] = None,
|
| 1225 |
+
proprio=None,
|
| 1226 |
+
proprio_projector=None,
|
| 1227 |
+
action_head:L1RegressionActionHead=None,
|
| 1228 |
+
noisy_action_projector=None,
|
| 1229 |
+
use_film: bool = False,
|
| 1230 |
+
**kwargs: str,
|
| 1231 |
+
) -> np.ndarray:
|
| 1232 |
+
|
| 1233 |
+
pixel_values = kwargs["pixel_values"]
|
| 1234 |
+
attention_mask = kwargs["attention_mask"]
|
| 1235 |
+
|
| 1236 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1237 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1238 |
+
|
| 1239 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1240 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1241 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1242 |
+
if cfg.preset:
|
| 1243 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1244 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1245 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1246 |
+
self.preset = True
|
| 1247 |
+
else:
|
| 1248 |
+
self.preset = False
|
| 1249 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1250 |
+
|
| 1251 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1252 |
+
|
| 1253 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1254 |
+
if use_proprio:
|
| 1255 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1256 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1257 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1258 |
+
)
|
| 1259 |
+
|
| 1260 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1261 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1262 |
+
# if use_proprio:
|
| 1263 |
+
# NUM_PATCHES += 1
|
| 1264 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1265 |
+
input_embeddings,
|
| 1266 |
+
None,
|
| 1267 |
+
projected_patch_embeddings,
|
| 1268 |
+
attention_mask,
|
| 1269 |
+
None, #推理不需要labels
|
| 1270 |
+
None, #推理不需要NUM_PATCHES
|
| 1271 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1272 |
+
action_head,
|
| 1273 |
+
)
|
| 1274 |
+
|
| 1275 |
+
# Unnormalize predicted actions
|
| 1276 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1277 |
+
|
| 1278 |
+
return actions, actions_hidden_states
|
| 1279 |
+
|
| 1280 |
+
|
| 1281 |
+
def mul_predict_action(
|
| 1282 |
+
self,
|
| 1283 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1284 |
+
unnorm_key: Optional[str] = None,
|
| 1285 |
+
proprio=None,
|
| 1286 |
+
proprio_projector=None,
|
| 1287 |
+
action_head:L1RegressionActionHead=None,
|
| 1288 |
+
noisy_action_projector=None,
|
| 1289 |
+
use_film: bool = False,
|
| 1290 |
+
**kwargs: str,
|
| 1291 |
+
) -> np.ndarray:
|
| 1292 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1293 |
+
if cfg.preset:
|
| 1294 |
+
self.preset = True
|
| 1295 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 1296 |
+
input_ids = torch.cat(
|
| 1297 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 1298 |
+
)
|
| 1299 |
+
else:
|
| 1300 |
+
self.preset = False
|
| 1301 |
+
|
| 1302 |
+
|
| 1303 |
+
pixel_values = kwargs["pixel_values"]
|
| 1304 |
+
attention_mask = kwargs["attention_mask"]
|
| 1305 |
+
|
| 1306 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1307 |
+
|
| 1308 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1309 |
+
|
| 1310 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1311 |
+
|
| 1312 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1313 |
+
if use_proprio:
|
| 1314 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1315 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1316 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1317 |
+
)
|
| 1318 |
+
|
| 1319 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1320 |
+
input_embeddings,
|
| 1321 |
+
None,
|
| 1322 |
+
projected_patch_embeddings,
|
| 1323 |
+
attention_mask,
|
| 1324 |
+
None, #推理不需要labels
|
| 1325 |
+
None, #推理不需要NUM_PATCHES
|
| 1326 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1327 |
+
action_head,
|
| 1328 |
+
)
|
| 1329 |
+
|
| 1330 |
+
# Unnormalize predicted actions
|
| 1331 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1332 |
+
|
| 1333 |
+
return actions, actions_hidden_states
|
| 1334 |
+
|
modeling_prismatic.py.back.20250403_110623
ADDED
|
@@ -0,0 +1,1321 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 957 |
+
import ipdb; ipdb.set_trace()
|
| 958 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 959 |
+
|
| 960 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 961 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 962 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 963 |
+
|
| 964 |
+
return normalized_actions, actions_hidden_states
|
| 965 |
+
|
| 966 |
+
def predict_action(
|
| 967 |
+
self,
|
| 968 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 969 |
+
unnorm_key: Optional[str] = None,
|
| 970 |
+
proprio=None,
|
| 971 |
+
proprio_projector=None,
|
| 972 |
+
action_head=None,
|
| 973 |
+
noisy_action_projector=None,
|
| 974 |
+
use_film: bool = False,
|
| 975 |
+
**kwargs: str,
|
| 976 |
+
) -> np.ndarray:
|
| 977 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 978 |
+
|
| 979 |
+
Args:
|
| 980 |
+
input_ids: Input token ids
|
| 981 |
+
unnorm_key: Key for unnormalization statistics
|
| 982 |
+
proprio: Proprioceptive features
|
| 983 |
+
proprio_projector: Projector for proprioceptive features
|
| 984 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 985 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 986 |
+
use_film: Whether to use FiLM conditioning
|
| 987 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 988 |
+
|
| 989 |
+
Returns:
|
| 990 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 991 |
+
"""
|
| 992 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 993 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 994 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 995 |
+
input_ids = torch.cat(
|
| 996 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 997 |
+
)
|
| 998 |
+
|
| 999 |
+
pixel_values = kwargs["pixel_values"]
|
| 1000 |
+
attention_mask = kwargs["attention_mask"]
|
| 1001 |
+
|
| 1002 |
+
# Create fake labels tensor (needed for action mask)
|
| 1003 |
+
labels = input_ids.clone()
|
| 1004 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1005 |
+
|
| 1006 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1007 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1008 |
+
|
| 1009 |
+
# Prepare inputs by adding necessary tokens
|
| 1010 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1011 |
+
|
| 1012 |
+
# Update labels tensor for action mask computation later
|
| 1013 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1014 |
+
|
| 1015 |
+
# Get input embeddings and action masks
|
| 1016 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1017 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1018 |
+
|
| 1019 |
+
# Extract language embeddings
|
| 1020 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1021 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1022 |
+
)
|
| 1023 |
+
|
| 1024 |
+
# Process vision features
|
| 1025 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1026 |
+
|
| 1027 |
+
# Add proprioceptive features if provided
|
| 1028 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1029 |
+
if use_proprio:
|
| 1030 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1031 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1032 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1033 |
+
)
|
| 1034 |
+
|
| 1035 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1036 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1037 |
+
|
| 1038 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1039 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1040 |
+
if use_proprio:
|
| 1041 |
+
NUM_PATCHES += 1
|
| 1042 |
+
|
| 1043 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1044 |
+
input_embeddings,
|
| 1045 |
+
all_actions_mask,
|
| 1046 |
+
projected_patch_embeddings,
|
| 1047 |
+
attention_mask,
|
| 1048 |
+
labels,
|
| 1049 |
+
NUM_PATCHES,
|
| 1050 |
+
NUM_PROMPT_TOKENS,
|
| 1051 |
+
action_head,
|
| 1052 |
+
)
|
| 1053 |
+
|
| 1054 |
+
# Unnormalize predicted actions
|
| 1055 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1056 |
+
|
| 1057 |
+
return actions, actions_hidden_states
|
| 1058 |
+
|
| 1059 |
+
@staticmethod
|
| 1060 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1061 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1062 |
+
if unnorm_key is None:
|
| 1063 |
+
assert len(norm_stats) == 1, (
|
| 1064 |
+
f"Your model was trained on more than one dataset, "
|
| 1065 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1066 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1067 |
+
)
|
| 1068 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1069 |
+
# norm states没有加载libero, 为什么?
|
| 1070 |
+
assert unnorm_key in norm_stats, (
|
| 1071 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1072 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1073 |
+
)
|
| 1074 |
+
return unnorm_key
|
| 1075 |
+
|
| 1076 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1077 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1078 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1079 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1080 |
+
|
| 1081 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1082 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1083 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1084 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1085 |
+
|
| 1086 |
+
|
| 1087 |
+
def lisa_forward(
|
| 1088 |
+
self,
|
| 1089 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1090 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1091 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1092 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1093 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1094 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1095 |
+
use_cache: Optional[bool] = None,
|
| 1096 |
+
output_attentions: Optional[bool] = None,
|
| 1097 |
+
output_hidden_states: Optional[bool] = None,
|
| 1098 |
+
output_projector_features: Optional[bool] = None,
|
| 1099 |
+
return_dict: Optional[bool] = None,
|
| 1100 |
+
proprio=None,
|
| 1101 |
+
proprio_projector=None,
|
| 1102 |
+
**kwargs
|
| 1103 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1104 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1105 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1106 |
+
output_hidden_states = (
|
| 1107 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1108 |
+
)
|
| 1109 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1110 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1111 |
+
|
| 1112 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1113 |
+
use_cache = use_cache and not self.training
|
| 1114 |
+
|
| 1115 |
+
# Instantiate Placeholder for Projector Features
|
| 1116 |
+
projected_patch_embeddings = None
|
| 1117 |
+
|
| 1118 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1119 |
+
if input_ids.shape[1] == 1:
|
| 1120 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1121 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1122 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1123 |
+
|
| 1124 |
+
language_model_output = self.language_model(
|
| 1125 |
+
input_ids=input_ids,
|
| 1126 |
+
attention_mask=None,
|
| 1127 |
+
position_ids=None,
|
| 1128 |
+
past_key_values=past_key_values,
|
| 1129 |
+
inputs_embeds=None,
|
| 1130 |
+
labels=None,
|
| 1131 |
+
use_cache=use_cache,
|
| 1132 |
+
output_attentions=output_attentions,
|
| 1133 |
+
output_hidden_states=output_hidden_states,
|
| 1134 |
+
return_dict=return_dict,
|
| 1135 |
+
)
|
| 1136 |
+
|
| 1137 |
+
# === Handle Unimodal Forward ===
|
| 1138 |
+
elif pixel_values is None:
|
| 1139 |
+
raise NotImplementedError
|
| 1140 |
+
|
| 1141 |
+
# === Handle Multimodal Forward ===
|
| 1142 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1143 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1144 |
+
|
| 1145 |
+
# Get input embeddings (from language model embeddings)
|
| 1146 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1147 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1148 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1149 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1150 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1151 |
+
# Get visual features
|
| 1152 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1153 |
+
|
| 1154 |
+
# Add proprioceptive state if provided
|
| 1155 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1156 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1157 |
+
)
|
| 1158 |
+
|
| 1159 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1160 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1161 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1162 |
+
|
| 1163 |
+
# Build multimodal embeddings & attention mask
|
| 1164 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1165 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1166 |
+
)
|
| 1167 |
+
|
| 1168 |
+
# Build labels for multimodal sequence if needed
|
| 1169 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1170 |
+
|
| 1171 |
+
# Dispatch to language model
|
| 1172 |
+
language_model_output = self.language_model(
|
| 1173 |
+
input_ids=None,
|
| 1174 |
+
attention_mask=multimodal_attention_mask,
|
| 1175 |
+
position_ids=None,
|
| 1176 |
+
past_key_values=None,
|
| 1177 |
+
inputs_embeds=multimodal_embeddings,
|
| 1178 |
+
labels=multimodal_labels,
|
| 1179 |
+
use_cache=use_cache,
|
| 1180 |
+
output_attentions=output_attentions,
|
| 1181 |
+
output_hidden_states=output_hidden_states,
|
| 1182 |
+
return_dict=return_dict,
|
| 1183 |
+
)
|
| 1184 |
+
|
| 1185 |
+
|
| 1186 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1187 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1188 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1189 |
+
|
| 1190 |
+
else:
|
| 1191 |
+
raise ValueError(
|
| 1192 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1193 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1194 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1195 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1196 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1197 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1198 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1199 |
+
f"=> `use_cache` = {use_cache}"
|
| 1200 |
+
)
|
| 1201 |
+
|
| 1202 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1203 |
+
if not return_dict:
|
| 1204 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1205 |
+
return *language_model_output, projected_patch_embeddings
|
| 1206 |
+
|
| 1207 |
+
return language_model_output
|
| 1208 |
+
|
| 1209 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1210 |
+
loss=language_model_output.loss,
|
| 1211 |
+
logits=language_model_output.logits,
|
| 1212 |
+
past_key_values=language_model_output.past_key_values,
|
| 1213 |
+
hidden_states=language_model_output.hidden_states,
|
| 1214 |
+
attentions=language_model_output.attentions,
|
| 1215 |
+
projector_features=projected_patch_embeddings,
|
| 1216 |
+
)
|
| 1217 |
+
|
| 1218 |
+
def lisa_predict_action(
|
| 1219 |
+
self,
|
| 1220 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1221 |
+
unnorm_key: Optional[str] = None,
|
| 1222 |
+
proprio=None,
|
| 1223 |
+
proprio_projector=None,
|
| 1224 |
+
action_head:L1RegressionActionHead=None,
|
| 1225 |
+
noisy_action_projector=None,
|
| 1226 |
+
use_film: bool = False,
|
| 1227 |
+
**kwargs: str,
|
| 1228 |
+
) -> np.ndarray:
|
| 1229 |
+
|
| 1230 |
+
pixel_values = kwargs["pixel_values"]
|
| 1231 |
+
attention_mask = kwargs["attention_mask"]
|
| 1232 |
+
|
| 1233 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1234 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1235 |
+
|
| 1236 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1237 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1238 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1239 |
+
if cfg.preset:
|
| 1240 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1241 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1242 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1243 |
+
self.preset = True
|
| 1244 |
+
else:
|
| 1245 |
+
self.preset = False
|
| 1246 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1247 |
+
|
| 1248 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1249 |
+
|
| 1250 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1251 |
+
if use_proprio:
|
| 1252 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1253 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1254 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1255 |
+
)
|
| 1256 |
+
|
| 1257 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1258 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1259 |
+
# if use_proprio:
|
| 1260 |
+
# NUM_PATCHES += 1
|
| 1261 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1262 |
+
input_embeddings,
|
| 1263 |
+
None,
|
| 1264 |
+
projected_patch_embeddings,
|
| 1265 |
+
attention_mask,
|
| 1266 |
+
None, #推理不需要labels
|
| 1267 |
+
None, #推理不需要NUM_PATCHES
|
| 1268 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1269 |
+
action_head,
|
| 1270 |
+
)
|
| 1271 |
+
|
| 1272 |
+
# Unnormalize predicted actions
|
| 1273 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1274 |
+
|
| 1275 |
+
return actions, actions_hidden_states
|
| 1276 |
+
|
| 1277 |
+
|
| 1278 |
+
def mul_predict_action(
|
| 1279 |
+
self,
|
| 1280 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1281 |
+
unnorm_key: Optional[str] = None,
|
| 1282 |
+
proprio=None,
|
| 1283 |
+
proprio_projector=None,
|
| 1284 |
+
action_head:L1RegressionActionHead=None,
|
| 1285 |
+
noisy_action_projector=None,
|
| 1286 |
+
use_film: bool = False,
|
| 1287 |
+
**kwargs: str,
|
| 1288 |
+
) -> np.ndarray:
|
| 1289 |
+
|
| 1290 |
+
pixel_values = kwargs["pixel_values"]
|
| 1291 |
+
attention_mask = kwargs["attention_mask"]
|
| 1292 |
+
|
| 1293 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1294 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1295 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1296 |
+
|
| 1297 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1298 |
+
|
| 1299 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1300 |
+
if use_proprio:
|
| 1301 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1302 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1303 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1304 |
+
)
|
| 1305 |
+
|
| 1306 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1307 |
+
input_embeddings,
|
| 1308 |
+
None,
|
| 1309 |
+
projected_patch_embeddings,
|
| 1310 |
+
attention_mask,
|
| 1311 |
+
None, #推理不需要labels
|
| 1312 |
+
None, #推理不需要NUM_PATCHES
|
| 1313 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1314 |
+
action_head,
|
| 1315 |
+
)
|
| 1316 |
+
|
| 1317 |
+
# Unnormalize predicted actions
|
| 1318 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1319 |
+
|
| 1320 |
+
return actions, actions_hidden_states
|
| 1321 |
+
|
modeling_prismatic.py.back.20250403_110730
ADDED
|
@@ -0,0 +1,1322 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 957 |
+
import ipdb; ipdb.set_trace()
|
| 958 |
+
# actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 959 |
+
actions_hidden_states = hidden_states[:, -1]
|
| 960 |
+
|
| 961 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 962 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 963 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 964 |
+
|
| 965 |
+
return normalized_actions, actions_hidden_states
|
| 966 |
+
|
| 967 |
+
def predict_action(
|
| 968 |
+
self,
|
| 969 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 970 |
+
unnorm_key: Optional[str] = None,
|
| 971 |
+
proprio=None,
|
| 972 |
+
proprio_projector=None,
|
| 973 |
+
action_head=None,
|
| 974 |
+
noisy_action_projector=None,
|
| 975 |
+
use_film: bool = False,
|
| 976 |
+
**kwargs: str,
|
| 977 |
+
) -> np.ndarray:
|
| 978 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 979 |
+
|
| 980 |
+
Args:
|
| 981 |
+
input_ids: Input token ids
|
| 982 |
+
unnorm_key: Key for unnormalization statistics
|
| 983 |
+
proprio: Proprioceptive features
|
| 984 |
+
proprio_projector: Projector for proprioceptive features
|
| 985 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 986 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 987 |
+
use_film: Whether to use FiLM conditioning
|
| 988 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 989 |
+
|
| 990 |
+
Returns:
|
| 991 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 992 |
+
"""
|
| 993 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 994 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 995 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 996 |
+
input_ids = torch.cat(
|
| 997 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 998 |
+
)
|
| 999 |
+
|
| 1000 |
+
pixel_values = kwargs["pixel_values"]
|
| 1001 |
+
attention_mask = kwargs["attention_mask"]
|
| 1002 |
+
|
| 1003 |
+
# Create fake labels tensor (needed for action mask)
|
| 1004 |
+
labels = input_ids.clone()
|
| 1005 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1006 |
+
|
| 1007 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1008 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1009 |
+
|
| 1010 |
+
# Prepare inputs by adding necessary tokens
|
| 1011 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1012 |
+
|
| 1013 |
+
# Update labels tensor for action mask computation later
|
| 1014 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1015 |
+
|
| 1016 |
+
# Get input embeddings and action masks
|
| 1017 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1018 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1019 |
+
|
| 1020 |
+
# Extract language embeddings
|
| 1021 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1022 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1023 |
+
)
|
| 1024 |
+
|
| 1025 |
+
# Process vision features
|
| 1026 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1027 |
+
|
| 1028 |
+
# Add proprioceptive features if provided
|
| 1029 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1030 |
+
if use_proprio:
|
| 1031 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1032 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1033 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1034 |
+
)
|
| 1035 |
+
|
| 1036 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1037 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1038 |
+
|
| 1039 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1040 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1041 |
+
if use_proprio:
|
| 1042 |
+
NUM_PATCHES += 1
|
| 1043 |
+
|
| 1044 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1045 |
+
input_embeddings,
|
| 1046 |
+
all_actions_mask,
|
| 1047 |
+
projected_patch_embeddings,
|
| 1048 |
+
attention_mask,
|
| 1049 |
+
labels,
|
| 1050 |
+
NUM_PATCHES,
|
| 1051 |
+
NUM_PROMPT_TOKENS,
|
| 1052 |
+
action_head,
|
| 1053 |
+
)
|
| 1054 |
+
|
| 1055 |
+
# Unnormalize predicted actions
|
| 1056 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1057 |
+
|
| 1058 |
+
return actions, actions_hidden_states
|
| 1059 |
+
|
| 1060 |
+
@staticmethod
|
| 1061 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1062 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1063 |
+
if unnorm_key is None:
|
| 1064 |
+
assert len(norm_stats) == 1, (
|
| 1065 |
+
f"Your model was trained on more than one dataset, "
|
| 1066 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1067 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1068 |
+
)
|
| 1069 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1070 |
+
# norm states没有加载libero, 为什么?
|
| 1071 |
+
assert unnorm_key in norm_stats, (
|
| 1072 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1073 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1074 |
+
)
|
| 1075 |
+
return unnorm_key
|
| 1076 |
+
|
| 1077 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1078 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1079 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1080 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1081 |
+
|
| 1082 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1083 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1084 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1085 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1086 |
+
|
| 1087 |
+
|
| 1088 |
+
def lisa_forward(
|
| 1089 |
+
self,
|
| 1090 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1091 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1092 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1093 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1094 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1095 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1096 |
+
use_cache: Optional[bool] = None,
|
| 1097 |
+
output_attentions: Optional[bool] = None,
|
| 1098 |
+
output_hidden_states: Optional[bool] = None,
|
| 1099 |
+
output_projector_features: Optional[bool] = None,
|
| 1100 |
+
return_dict: Optional[bool] = None,
|
| 1101 |
+
proprio=None,
|
| 1102 |
+
proprio_projector=None,
|
| 1103 |
+
**kwargs
|
| 1104 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1105 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1106 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1107 |
+
output_hidden_states = (
|
| 1108 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1109 |
+
)
|
| 1110 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1111 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1112 |
+
|
| 1113 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1114 |
+
use_cache = use_cache and not self.training
|
| 1115 |
+
|
| 1116 |
+
# Instantiate Placeholder for Projector Features
|
| 1117 |
+
projected_patch_embeddings = None
|
| 1118 |
+
|
| 1119 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1120 |
+
if input_ids.shape[1] == 1:
|
| 1121 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1122 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1123 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1124 |
+
|
| 1125 |
+
language_model_output = self.language_model(
|
| 1126 |
+
input_ids=input_ids,
|
| 1127 |
+
attention_mask=None,
|
| 1128 |
+
position_ids=None,
|
| 1129 |
+
past_key_values=past_key_values,
|
| 1130 |
+
inputs_embeds=None,
|
| 1131 |
+
labels=None,
|
| 1132 |
+
use_cache=use_cache,
|
| 1133 |
+
output_attentions=output_attentions,
|
| 1134 |
+
output_hidden_states=output_hidden_states,
|
| 1135 |
+
return_dict=return_dict,
|
| 1136 |
+
)
|
| 1137 |
+
|
| 1138 |
+
# === Handle Unimodal Forward ===
|
| 1139 |
+
elif pixel_values is None:
|
| 1140 |
+
raise NotImplementedError
|
| 1141 |
+
|
| 1142 |
+
# === Handle Multimodal Forward ===
|
| 1143 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1144 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1145 |
+
|
| 1146 |
+
# Get input embeddings (from language model embeddings)
|
| 1147 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1148 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1149 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1150 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1151 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1152 |
+
# Get visual features
|
| 1153 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1154 |
+
|
| 1155 |
+
# Add proprioceptive state if provided
|
| 1156 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1157 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1158 |
+
)
|
| 1159 |
+
|
| 1160 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1161 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1162 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1163 |
+
|
| 1164 |
+
# Build multimodal embeddings & attention mask
|
| 1165 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1166 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1167 |
+
)
|
| 1168 |
+
|
| 1169 |
+
# Build labels for multimodal sequence if needed
|
| 1170 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1171 |
+
|
| 1172 |
+
# Dispatch to language model
|
| 1173 |
+
language_model_output = self.language_model(
|
| 1174 |
+
input_ids=None,
|
| 1175 |
+
attention_mask=multimodal_attention_mask,
|
| 1176 |
+
position_ids=None,
|
| 1177 |
+
past_key_values=None,
|
| 1178 |
+
inputs_embeds=multimodal_embeddings,
|
| 1179 |
+
labels=multimodal_labels,
|
| 1180 |
+
use_cache=use_cache,
|
| 1181 |
+
output_attentions=output_attentions,
|
| 1182 |
+
output_hidden_states=output_hidden_states,
|
| 1183 |
+
return_dict=return_dict,
|
| 1184 |
+
)
|
| 1185 |
+
|
| 1186 |
+
|
| 1187 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1188 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1189 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1190 |
+
|
| 1191 |
+
else:
|
| 1192 |
+
raise ValueError(
|
| 1193 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1194 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1195 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1196 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1197 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1198 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1199 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1200 |
+
f"=> `use_cache` = {use_cache}"
|
| 1201 |
+
)
|
| 1202 |
+
|
| 1203 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1204 |
+
if not return_dict:
|
| 1205 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1206 |
+
return *language_model_output, projected_patch_embeddings
|
| 1207 |
+
|
| 1208 |
+
return language_model_output
|
| 1209 |
+
|
| 1210 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1211 |
+
loss=language_model_output.loss,
|
| 1212 |
+
logits=language_model_output.logits,
|
| 1213 |
+
past_key_values=language_model_output.past_key_values,
|
| 1214 |
+
hidden_states=language_model_output.hidden_states,
|
| 1215 |
+
attentions=language_model_output.attentions,
|
| 1216 |
+
projector_features=projected_patch_embeddings,
|
| 1217 |
+
)
|
| 1218 |
+
|
| 1219 |
+
def lisa_predict_action(
|
| 1220 |
+
self,
|
| 1221 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1222 |
+
unnorm_key: Optional[str] = None,
|
| 1223 |
+
proprio=None,
|
| 1224 |
+
proprio_projector=None,
|
| 1225 |
+
action_head:L1RegressionActionHead=None,
|
| 1226 |
+
noisy_action_projector=None,
|
| 1227 |
+
use_film: bool = False,
|
| 1228 |
+
**kwargs: str,
|
| 1229 |
+
) -> np.ndarray:
|
| 1230 |
+
|
| 1231 |
+
pixel_values = kwargs["pixel_values"]
|
| 1232 |
+
attention_mask = kwargs["attention_mask"]
|
| 1233 |
+
|
| 1234 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1235 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1236 |
+
|
| 1237 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1238 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1239 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1240 |
+
if cfg.preset:
|
| 1241 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1242 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1243 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1244 |
+
self.preset = True
|
| 1245 |
+
else:
|
| 1246 |
+
self.preset = False
|
| 1247 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1248 |
+
|
| 1249 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1250 |
+
|
| 1251 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1252 |
+
if use_proprio:
|
| 1253 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1254 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1255 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1256 |
+
)
|
| 1257 |
+
|
| 1258 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1259 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1260 |
+
# if use_proprio:
|
| 1261 |
+
# NUM_PATCHES += 1
|
| 1262 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1263 |
+
input_embeddings,
|
| 1264 |
+
None,
|
| 1265 |
+
projected_patch_embeddings,
|
| 1266 |
+
attention_mask,
|
| 1267 |
+
None, #推理不需要labels
|
| 1268 |
+
None, #推理不需要NUM_PATCHES
|
| 1269 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1270 |
+
action_head,
|
| 1271 |
+
)
|
| 1272 |
+
|
| 1273 |
+
# Unnormalize predicted actions
|
| 1274 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1275 |
+
|
| 1276 |
+
return actions, actions_hidden_states
|
| 1277 |
+
|
| 1278 |
+
|
| 1279 |
+
def mul_predict_action(
|
| 1280 |
+
self,
|
| 1281 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1282 |
+
unnorm_key: Optional[str] = None,
|
| 1283 |
+
proprio=None,
|
| 1284 |
+
proprio_projector=None,
|
| 1285 |
+
action_head:L1RegressionActionHead=None,
|
| 1286 |
+
noisy_action_projector=None,
|
| 1287 |
+
use_film: bool = False,
|
| 1288 |
+
**kwargs: str,
|
| 1289 |
+
) -> np.ndarray:
|
| 1290 |
+
|
| 1291 |
+
pixel_values = kwargs["pixel_values"]
|
| 1292 |
+
attention_mask = kwargs["attention_mask"]
|
| 1293 |
+
|
| 1294 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1295 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1296 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1297 |
+
|
| 1298 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1299 |
+
|
| 1300 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1301 |
+
if use_proprio:
|
| 1302 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1303 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1304 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1305 |
+
)
|
| 1306 |
+
|
| 1307 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1308 |
+
input_embeddings,
|
| 1309 |
+
None,
|
| 1310 |
+
projected_patch_embeddings,
|
| 1311 |
+
attention_mask,
|
| 1312 |
+
None, #推理不需要labels
|
| 1313 |
+
None, #推理不需要NUM_PATCHES
|
| 1314 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1315 |
+
action_head,
|
| 1316 |
+
)
|
| 1317 |
+
|
| 1318 |
+
# Unnormalize predicted actions
|
| 1319 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1320 |
+
|
| 1321 |
+
return actions, actions_hidden_states
|
| 1322 |
+
|
modeling_prismatic.py.back.20250403_111133
ADDED
|
@@ -0,0 +1,1322 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 957 |
+
import ipdb; ipdb.set_trace()
|
| 958 |
+
# actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 959 |
+
actions_hidden_states = hidden_states[:, -1]
|
| 960 |
+
|
| 961 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 962 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 963 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 964 |
+
|
| 965 |
+
return normalized_actions, actions_hidden_states
|
| 966 |
+
|
| 967 |
+
def predict_action(
|
| 968 |
+
self,
|
| 969 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 970 |
+
unnorm_key: Optional[str] = None,
|
| 971 |
+
proprio=None,
|
| 972 |
+
proprio_projector=None,
|
| 973 |
+
action_head=None,
|
| 974 |
+
noisy_action_projector=None,
|
| 975 |
+
use_film: bool = False,
|
| 976 |
+
**kwargs: str,
|
| 977 |
+
) -> np.ndarray:
|
| 978 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 979 |
+
|
| 980 |
+
Args:
|
| 981 |
+
input_ids: Input token ids
|
| 982 |
+
unnorm_key: Key for unnormalization statistics
|
| 983 |
+
proprio: Proprioceptive features
|
| 984 |
+
proprio_projector: Projector for proprioceptive features
|
| 985 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 986 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 987 |
+
use_film: Whether to use FiLM conditioning
|
| 988 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 989 |
+
|
| 990 |
+
Returns:
|
| 991 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 992 |
+
"""
|
| 993 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 994 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 995 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 996 |
+
input_ids = torch.cat(
|
| 997 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 998 |
+
)
|
| 999 |
+
|
| 1000 |
+
pixel_values = kwargs["pixel_values"]
|
| 1001 |
+
attention_mask = kwargs["attention_mask"]
|
| 1002 |
+
|
| 1003 |
+
# Create fake labels tensor (needed for action mask)
|
| 1004 |
+
labels = input_ids.clone()
|
| 1005 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1006 |
+
|
| 1007 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1008 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1009 |
+
|
| 1010 |
+
# Prepare inputs by adding necessary tokens
|
| 1011 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1012 |
+
|
| 1013 |
+
# Update labels tensor for action mask computation later
|
| 1014 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1015 |
+
|
| 1016 |
+
# Get input embeddings and action masks
|
| 1017 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1018 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1019 |
+
|
| 1020 |
+
# Extract language embeddings
|
| 1021 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1022 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1023 |
+
)
|
| 1024 |
+
|
| 1025 |
+
# Process vision features
|
| 1026 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1027 |
+
|
| 1028 |
+
# Add proprioceptive features if provided
|
| 1029 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1030 |
+
if use_proprio:
|
| 1031 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1032 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1033 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1034 |
+
)
|
| 1035 |
+
|
| 1036 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1037 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1038 |
+
|
| 1039 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1040 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1041 |
+
if use_proprio:
|
| 1042 |
+
NUM_PATCHES += 1
|
| 1043 |
+
|
| 1044 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1045 |
+
input_embeddings,
|
| 1046 |
+
all_actions_mask,
|
| 1047 |
+
projected_patch_embeddings,
|
| 1048 |
+
attention_mask,
|
| 1049 |
+
labels,
|
| 1050 |
+
NUM_PATCHES,
|
| 1051 |
+
NUM_PROMPT_TOKENS,
|
| 1052 |
+
action_head,
|
| 1053 |
+
)
|
| 1054 |
+
|
| 1055 |
+
# Unnormalize predicted actions
|
| 1056 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1057 |
+
|
| 1058 |
+
return actions, actions_hidden_states
|
| 1059 |
+
|
| 1060 |
+
@staticmethod
|
| 1061 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1062 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1063 |
+
if unnorm_key is None:
|
| 1064 |
+
assert len(norm_stats) == 1, (
|
| 1065 |
+
f"Your model was trained on more than one dataset, "
|
| 1066 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1067 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1068 |
+
)
|
| 1069 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1070 |
+
# norm states没有加载libero, 为什么?
|
| 1071 |
+
assert unnorm_key in norm_stats, (
|
| 1072 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1073 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1074 |
+
)
|
| 1075 |
+
return unnorm_key
|
| 1076 |
+
|
| 1077 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1078 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1079 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1080 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1081 |
+
|
| 1082 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1083 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1084 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1085 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1086 |
+
|
| 1087 |
+
|
| 1088 |
+
def lisa_forward(
|
| 1089 |
+
self,
|
| 1090 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1091 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1092 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1093 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1094 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1095 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1096 |
+
use_cache: Optional[bool] = None,
|
| 1097 |
+
output_attentions: Optional[bool] = None,
|
| 1098 |
+
output_hidden_states: Optional[bool] = None,
|
| 1099 |
+
output_projector_features: Optional[bool] = None,
|
| 1100 |
+
return_dict: Optional[bool] = None,
|
| 1101 |
+
proprio=None,
|
| 1102 |
+
proprio_projector=None,
|
| 1103 |
+
**kwargs
|
| 1104 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1105 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1106 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1107 |
+
output_hidden_states = (
|
| 1108 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1109 |
+
)
|
| 1110 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1111 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1112 |
+
|
| 1113 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1114 |
+
use_cache = use_cache and not self.training
|
| 1115 |
+
|
| 1116 |
+
# Instantiate Placeholder for Projector Features
|
| 1117 |
+
projected_patch_embeddings = None
|
| 1118 |
+
|
| 1119 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1120 |
+
if input_ids.shape[1] == 1:
|
| 1121 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1122 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1123 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1124 |
+
|
| 1125 |
+
language_model_output = self.language_model(
|
| 1126 |
+
input_ids=input_ids,
|
| 1127 |
+
attention_mask=None,
|
| 1128 |
+
position_ids=None,
|
| 1129 |
+
past_key_values=past_key_values,
|
| 1130 |
+
inputs_embeds=None,
|
| 1131 |
+
labels=None,
|
| 1132 |
+
use_cache=use_cache,
|
| 1133 |
+
output_attentions=output_attentions,
|
| 1134 |
+
output_hidden_states=output_hidden_states,
|
| 1135 |
+
return_dict=return_dict,
|
| 1136 |
+
)
|
| 1137 |
+
|
| 1138 |
+
# === Handle Unimodal Forward ===
|
| 1139 |
+
elif pixel_values is None:
|
| 1140 |
+
raise NotImplementedError
|
| 1141 |
+
|
| 1142 |
+
# === Handle Multimodal Forward ===
|
| 1143 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1144 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1145 |
+
|
| 1146 |
+
# Get input embeddings (from language model embeddings)
|
| 1147 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1148 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1149 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1150 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1151 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1152 |
+
# Get visual features
|
| 1153 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1154 |
+
|
| 1155 |
+
# Add proprioceptive state if provided
|
| 1156 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1157 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1158 |
+
)
|
| 1159 |
+
|
| 1160 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1161 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1162 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1163 |
+
|
| 1164 |
+
# Build multimodal embeddings & attention mask
|
| 1165 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1166 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1167 |
+
)
|
| 1168 |
+
|
| 1169 |
+
# Build labels for multimodal sequence if needed
|
| 1170 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1171 |
+
|
| 1172 |
+
# Dispatch to language model
|
| 1173 |
+
language_model_output = self.language_model(
|
| 1174 |
+
input_ids=None,
|
| 1175 |
+
attention_mask=multimodal_attention_mask,
|
| 1176 |
+
position_ids=None,
|
| 1177 |
+
past_key_values=None,
|
| 1178 |
+
inputs_embeds=multimodal_embeddings,
|
| 1179 |
+
labels=multimodal_labels,
|
| 1180 |
+
use_cache=use_cache,
|
| 1181 |
+
output_attentions=output_attentions,
|
| 1182 |
+
output_hidden_states=output_hidden_states,
|
| 1183 |
+
return_dict=return_dict,
|
| 1184 |
+
)
|
| 1185 |
+
|
| 1186 |
+
|
| 1187 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1188 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1189 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1190 |
+
|
| 1191 |
+
else:
|
| 1192 |
+
raise ValueError(
|
| 1193 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1194 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1195 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1196 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1197 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1198 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1199 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1200 |
+
f"=> `use_cache` = {use_cache}"
|
| 1201 |
+
)
|
| 1202 |
+
|
| 1203 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1204 |
+
if not return_dict:
|
| 1205 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1206 |
+
return *language_model_output, projected_patch_embeddings
|
| 1207 |
+
|
| 1208 |
+
return language_model_output
|
| 1209 |
+
|
| 1210 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1211 |
+
loss=language_model_output.loss,
|
| 1212 |
+
logits=language_model_output.logits,
|
| 1213 |
+
past_key_values=language_model_output.past_key_values,
|
| 1214 |
+
hidden_states=language_model_output.hidden_states,
|
| 1215 |
+
attentions=language_model_output.attentions,
|
| 1216 |
+
projector_features=projected_patch_embeddings,
|
| 1217 |
+
)
|
| 1218 |
+
|
| 1219 |
+
def lisa_predict_action(
|
| 1220 |
+
self,
|
| 1221 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1222 |
+
unnorm_key: Optional[str] = None,
|
| 1223 |
+
proprio=None,
|
| 1224 |
+
proprio_projector=None,
|
| 1225 |
+
action_head:L1RegressionActionHead=None,
|
| 1226 |
+
noisy_action_projector=None,
|
| 1227 |
+
use_film: bool = False,
|
| 1228 |
+
**kwargs: str,
|
| 1229 |
+
) -> np.ndarray:
|
| 1230 |
+
|
| 1231 |
+
pixel_values = kwargs["pixel_values"]
|
| 1232 |
+
attention_mask = kwargs["attention_mask"]
|
| 1233 |
+
|
| 1234 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1235 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1236 |
+
|
| 1237 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1238 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1239 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1240 |
+
if cfg.preset:
|
| 1241 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1242 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1243 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1244 |
+
self.preset = True
|
| 1245 |
+
else:
|
| 1246 |
+
self.preset = False
|
| 1247 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1248 |
+
|
| 1249 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1250 |
+
|
| 1251 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1252 |
+
if use_proprio:
|
| 1253 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1254 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1255 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1256 |
+
)
|
| 1257 |
+
|
| 1258 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1259 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1260 |
+
# if use_proprio:
|
| 1261 |
+
# NUM_PATCHES += 1
|
| 1262 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1263 |
+
input_embeddings,
|
| 1264 |
+
None,
|
| 1265 |
+
projected_patch_embeddings,
|
| 1266 |
+
attention_mask,
|
| 1267 |
+
None, #推理不需要labels
|
| 1268 |
+
None, #推理不需要NUM_PATCHES
|
| 1269 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1270 |
+
action_head,
|
| 1271 |
+
)
|
| 1272 |
+
|
| 1273 |
+
# Unnormalize predicted actions
|
| 1274 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1275 |
+
|
| 1276 |
+
return actions, actions_hidden_states
|
| 1277 |
+
|
| 1278 |
+
|
| 1279 |
+
def mul_predict_action(
|
| 1280 |
+
self,
|
| 1281 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1282 |
+
unnorm_key: Optional[str] = None,
|
| 1283 |
+
proprio=None,
|
| 1284 |
+
proprio_projector=None,
|
| 1285 |
+
action_head:L1RegressionActionHead=None,
|
| 1286 |
+
noisy_action_projector=None,
|
| 1287 |
+
use_film: bool = False,
|
| 1288 |
+
**kwargs: str,
|
| 1289 |
+
) -> np.ndarray:
|
| 1290 |
+
|
| 1291 |
+
pixel_values = kwargs["pixel_values"]
|
| 1292 |
+
attention_mask = kwargs["attention_mask"]
|
| 1293 |
+
|
| 1294 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1295 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1296 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1297 |
+
|
| 1298 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1299 |
+
|
| 1300 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1301 |
+
if use_proprio:
|
| 1302 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1303 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1304 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1305 |
+
)
|
| 1306 |
+
|
| 1307 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1308 |
+
input_embeddings,
|
| 1309 |
+
None,
|
| 1310 |
+
projected_patch_embeddings,
|
| 1311 |
+
attention_mask,
|
| 1312 |
+
None, #推理不需要labels
|
| 1313 |
+
None, #推理不需要NUM_PATCHES
|
| 1314 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1315 |
+
action_head,
|
| 1316 |
+
)
|
| 1317 |
+
|
| 1318 |
+
# Unnormalize predicted actions
|
| 1319 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1320 |
+
|
| 1321 |
+
return actions, actions_hidden_states
|
| 1322 |
+
|
modeling_prismatic.py.back.20250403_111217
ADDED
|
@@ -0,0 +1,1323 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
# 0 , is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 957 |
+
actions_hidden_states = language_model_output.hidden_states[0][-1]
|
| 958 |
+
# import ipdb; ipdb.set_trace()
|
| 959 |
+
# actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 960 |
+
actions_hidden_states = hidden_states[:, -1]
|
| 961 |
+
|
| 962 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 963 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 964 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 965 |
+
|
| 966 |
+
return normalized_actions, actions_hidden_states
|
| 967 |
+
|
| 968 |
+
def predict_action(
|
| 969 |
+
self,
|
| 970 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 971 |
+
unnorm_key: Optional[str] = None,
|
| 972 |
+
proprio=None,
|
| 973 |
+
proprio_projector=None,
|
| 974 |
+
action_head=None,
|
| 975 |
+
noisy_action_projector=None,
|
| 976 |
+
use_film: bool = False,
|
| 977 |
+
**kwargs: str,
|
| 978 |
+
) -> np.ndarray:
|
| 979 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 980 |
+
|
| 981 |
+
Args:
|
| 982 |
+
input_ids: Input token ids
|
| 983 |
+
unnorm_key: Key for unnormalization statistics
|
| 984 |
+
proprio: Proprioceptive features
|
| 985 |
+
proprio_projector: Projector for proprioceptive features
|
| 986 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 987 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 988 |
+
use_film: Whether to use FiLM conditioning
|
| 989 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 990 |
+
|
| 991 |
+
Returns:
|
| 992 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 993 |
+
"""
|
| 994 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 995 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 996 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 997 |
+
input_ids = torch.cat(
|
| 998 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 999 |
+
)
|
| 1000 |
+
|
| 1001 |
+
pixel_values = kwargs["pixel_values"]
|
| 1002 |
+
attention_mask = kwargs["attention_mask"]
|
| 1003 |
+
|
| 1004 |
+
# Create fake labels tensor (needed for action mask)
|
| 1005 |
+
labels = input_ids.clone()
|
| 1006 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1007 |
+
|
| 1008 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1009 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1010 |
+
|
| 1011 |
+
# Prepare inputs by adding necessary tokens
|
| 1012 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1013 |
+
|
| 1014 |
+
# Update labels tensor for action mask computation later
|
| 1015 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1016 |
+
|
| 1017 |
+
# Get input embeddings and action masks
|
| 1018 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1019 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1020 |
+
|
| 1021 |
+
# Extract language embeddings
|
| 1022 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1023 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1024 |
+
)
|
| 1025 |
+
|
| 1026 |
+
# Process vision features
|
| 1027 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1028 |
+
|
| 1029 |
+
# Add proprioceptive features if provided
|
| 1030 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1031 |
+
if use_proprio:
|
| 1032 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1033 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1034 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1035 |
+
)
|
| 1036 |
+
|
| 1037 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1038 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1039 |
+
|
| 1040 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1041 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1042 |
+
if use_proprio:
|
| 1043 |
+
NUM_PATCHES += 1
|
| 1044 |
+
|
| 1045 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1046 |
+
input_embeddings,
|
| 1047 |
+
all_actions_mask,
|
| 1048 |
+
projected_patch_embeddings,
|
| 1049 |
+
attention_mask,
|
| 1050 |
+
labels,
|
| 1051 |
+
NUM_PATCHES,
|
| 1052 |
+
NUM_PROMPT_TOKENS,
|
| 1053 |
+
action_head,
|
| 1054 |
+
)
|
| 1055 |
+
|
| 1056 |
+
# Unnormalize predicted actions
|
| 1057 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1058 |
+
|
| 1059 |
+
return actions, actions_hidden_states
|
| 1060 |
+
|
| 1061 |
+
@staticmethod
|
| 1062 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1063 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1064 |
+
if unnorm_key is None:
|
| 1065 |
+
assert len(norm_stats) == 1, (
|
| 1066 |
+
f"Your model was trained on more than one dataset, "
|
| 1067 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1068 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1069 |
+
)
|
| 1070 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1071 |
+
# norm states没有加载libero, 为什么?
|
| 1072 |
+
assert unnorm_key in norm_stats, (
|
| 1073 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1074 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1075 |
+
)
|
| 1076 |
+
return unnorm_key
|
| 1077 |
+
|
| 1078 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1079 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1080 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1081 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1082 |
+
|
| 1083 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1084 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1085 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1086 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1087 |
+
|
| 1088 |
+
|
| 1089 |
+
def lisa_forward(
|
| 1090 |
+
self,
|
| 1091 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1092 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1093 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1094 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1095 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1096 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1097 |
+
use_cache: Optional[bool] = None,
|
| 1098 |
+
output_attentions: Optional[bool] = None,
|
| 1099 |
+
output_hidden_states: Optional[bool] = None,
|
| 1100 |
+
output_projector_features: Optional[bool] = None,
|
| 1101 |
+
return_dict: Optional[bool] = None,
|
| 1102 |
+
proprio=None,
|
| 1103 |
+
proprio_projector=None,
|
| 1104 |
+
**kwargs
|
| 1105 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1106 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1107 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1108 |
+
output_hidden_states = (
|
| 1109 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1110 |
+
)
|
| 1111 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1112 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1113 |
+
|
| 1114 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1115 |
+
use_cache = use_cache and not self.training
|
| 1116 |
+
|
| 1117 |
+
# Instantiate Placeholder for Projector Features
|
| 1118 |
+
projected_patch_embeddings = None
|
| 1119 |
+
|
| 1120 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1121 |
+
if input_ids.shape[1] == 1:
|
| 1122 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1123 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1124 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1125 |
+
|
| 1126 |
+
language_model_output = self.language_model(
|
| 1127 |
+
input_ids=input_ids,
|
| 1128 |
+
attention_mask=None,
|
| 1129 |
+
position_ids=None,
|
| 1130 |
+
past_key_values=past_key_values,
|
| 1131 |
+
inputs_embeds=None,
|
| 1132 |
+
labels=None,
|
| 1133 |
+
use_cache=use_cache,
|
| 1134 |
+
output_attentions=output_attentions,
|
| 1135 |
+
output_hidden_states=output_hidden_states,
|
| 1136 |
+
return_dict=return_dict,
|
| 1137 |
+
)
|
| 1138 |
+
|
| 1139 |
+
# === Handle Unimodal Forward ===
|
| 1140 |
+
elif pixel_values is None:
|
| 1141 |
+
raise NotImplementedError
|
| 1142 |
+
|
| 1143 |
+
# === Handle Multimodal Forward ===
|
| 1144 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1145 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1146 |
+
|
| 1147 |
+
# Get input embeddings (from language model embeddings)
|
| 1148 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1149 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1150 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1151 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1152 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1153 |
+
# Get visual features
|
| 1154 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1155 |
+
|
| 1156 |
+
# Add proprioceptive state if provided
|
| 1157 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1158 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1159 |
+
)
|
| 1160 |
+
|
| 1161 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1162 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1163 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1164 |
+
|
| 1165 |
+
# Build multimodal embeddings & attention mask
|
| 1166 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1167 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1168 |
+
)
|
| 1169 |
+
|
| 1170 |
+
# Build labels for multimodal sequence if needed
|
| 1171 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1172 |
+
|
| 1173 |
+
# Dispatch to language model
|
| 1174 |
+
language_model_output = self.language_model(
|
| 1175 |
+
input_ids=None,
|
| 1176 |
+
attention_mask=multimodal_attention_mask,
|
| 1177 |
+
position_ids=None,
|
| 1178 |
+
past_key_values=None,
|
| 1179 |
+
inputs_embeds=multimodal_embeddings,
|
| 1180 |
+
labels=multimodal_labels,
|
| 1181 |
+
use_cache=use_cache,
|
| 1182 |
+
output_attentions=output_attentions,
|
| 1183 |
+
output_hidden_states=output_hidden_states,
|
| 1184 |
+
return_dict=return_dict,
|
| 1185 |
+
)
|
| 1186 |
+
|
| 1187 |
+
|
| 1188 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1189 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1190 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1191 |
+
|
| 1192 |
+
else:
|
| 1193 |
+
raise ValueError(
|
| 1194 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1195 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1196 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1197 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1198 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1199 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1200 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1201 |
+
f"=> `use_cache` = {use_cache}"
|
| 1202 |
+
)
|
| 1203 |
+
|
| 1204 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1205 |
+
if not return_dict:
|
| 1206 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1207 |
+
return *language_model_output, projected_patch_embeddings
|
| 1208 |
+
|
| 1209 |
+
return language_model_output
|
| 1210 |
+
|
| 1211 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1212 |
+
loss=language_model_output.loss,
|
| 1213 |
+
logits=language_model_output.logits,
|
| 1214 |
+
past_key_values=language_model_output.past_key_values,
|
| 1215 |
+
hidden_states=language_model_output.hidden_states,
|
| 1216 |
+
attentions=language_model_output.attentions,
|
| 1217 |
+
projector_features=projected_patch_embeddings,
|
| 1218 |
+
)
|
| 1219 |
+
|
| 1220 |
+
def lisa_predict_action(
|
| 1221 |
+
self,
|
| 1222 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1223 |
+
unnorm_key: Optional[str] = None,
|
| 1224 |
+
proprio=None,
|
| 1225 |
+
proprio_projector=None,
|
| 1226 |
+
action_head:L1RegressionActionHead=None,
|
| 1227 |
+
noisy_action_projector=None,
|
| 1228 |
+
use_film: bool = False,
|
| 1229 |
+
**kwargs: str,
|
| 1230 |
+
) -> np.ndarray:
|
| 1231 |
+
|
| 1232 |
+
pixel_values = kwargs["pixel_values"]
|
| 1233 |
+
attention_mask = kwargs["attention_mask"]
|
| 1234 |
+
|
| 1235 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1236 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1237 |
+
|
| 1238 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1239 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1240 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1241 |
+
if cfg.preset:
|
| 1242 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1243 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1244 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1245 |
+
self.preset = True
|
| 1246 |
+
else:
|
| 1247 |
+
self.preset = False
|
| 1248 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1249 |
+
|
| 1250 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1251 |
+
|
| 1252 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1253 |
+
if use_proprio:
|
| 1254 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1255 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1256 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1257 |
+
)
|
| 1258 |
+
|
| 1259 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1260 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1261 |
+
# if use_proprio:
|
| 1262 |
+
# NUM_PATCHES += 1
|
| 1263 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1264 |
+
input_embeddings,
|
| 1265 |
+
None,
|
| 1266 |
+
projected_patch_embeddings,
|
| 1267 |
+
attention_mask,
|
| 1268 |
+
None, #推理不需要labels
|
| 1269 |
+
None, #推理不需要NUM_PATCHES
|
| 1270 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1271 |
+
action_head,
|
| 1272 |
+
)
|
| 1273 |
+
|
| 1274 |
+
# Unnormalize predicted actions
|
| 1275 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1276 |
+
|
| 1277 |
+
return actions, actions_hidden_states
|
| 1278 |
+
|
| 1279 |
+
|
| 1280 |
+
def mul_predict_action(
|
| 1281 |
+
self,
|
| 1282 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1283 |
+
unnorm_key: Optional[str] = None,
|
| 1284 |
+
proprio=None,
|
| 1285 |
+
proprio_projector=None,
|
| 1286 |
+
action_head:L1RegressionActionHead=None,
|
| 1287 |
+
noisy_action_projector=None,
|
| 1288 |
+
use_film: bool = False,
|
| 1289 |
+
**kwargs: str,
|
| 1290 |
+
) -> np.ndarray:
|
| 1291 |
+
|
| 1292 |
+
pixel_values = kwargs["pixel_values"]
|
| 1293 |
+
attention_mask = kwargs["attention_mask"]
|
| 1294 |
+
|
| 1295 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1296 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1297 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1298 |
+
|
| 1299 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1300 |
+
|
| 1301 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1302 |
+
if use_proprio:
|
| 1303 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1304 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1305 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1306 |
+
)
|
| 1307 |
+
|
| 1308 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1309 |
+
input_embeddings,
|
| 1310 |
+
None,
|
| 1311 |
+
projected_patch_embeddings,
|
| 1312 |
+
attention_mask,
|
| 1313 |
+
None, #推理不需要labels
|
| 1314 |
+
None, #推理不需要NUM_PATCHES
|
| 1315 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1316 |
+
action_head,
|
| 1317 |
+
)
|
| 1318 |
+
|
| 1319 |
+
# Unnormalize predicted actions
|
| 1320 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1321 |
+
|
| 1322 |
+
return actions, actions_hidden_states
|
| 1323 |
+
|
modeling_prismatic.py.back.20250403_111705
ADDED
|
@@ -0,0 +1,1323 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
# 0 , is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 957 |
+
actions_hidden_states = language_model_output.hidden_states[0][-1] # 第一个是29871, 应该acc低, 第二个act, 应该acc高.
|
| 958 |
+
# import ipdb; ipdb.set_trace()
|
| 959 |
+
# actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 960 |
+
actions_hidden_states = actions_hidden_states[:, -1]
|
| 961 |
+
|
| 962 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 963 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 964 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 965 |
+
|
| 966 |
+
return normalized_actions, actions_hidden_states
|
| 967 |
+
|
| 968 |
+
def predict_action(
|
| 969 |
+
self,
|
| 970 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 971 |
+
unnorm_key: Optional[str] = None,
|
| 972 |
+
proprio=None,
|
| 973 |
+
proprio_projector=None,
|
| 974 |
+
action_head=None,
|
| 975 |
+
noisy_action_projector=None,
|
| 976 |
+
use_film: bool = False,
|
| 977 |
+
**kwargs: str,
|
| 978 |
+
) -> np.ndarray:
|
| 979 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 980 |
+
|
| 981 |
+
Args:
|
| 982 |
+
input_ids: Input token ids
|
| 983 |
+
unnorm_key: Key for unnormalization statistics
|
| 984 |
+
proprio: Proprioceptive features
|
| 985 |
+
proprio_projector: Projector for proprioceptive features
|
| 986 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 987 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 988 |
+
use_film: Whether to use FiLM conditioning
|
| 989 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 990 |
+
|
| 991 |
+
Returns:
|
| 992 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 993 |
+
"""
|
| 994 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 995 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 996 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 997 |
+
input_ids = torch.cat(
|
| 998 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 999 |
+
)
|
| 1000 |
+
|
| 1001 |
+
pixel_values = kwargs["pixel_values"]
|
| 1002 |
+
attention_mask = kwargs["attention_mask"]
|
| 1003 |
+
|
| 1004 |
+
# Create fake labels tensor (needed for action mask)
|
| 1005 |
+
labels = input_ids.clone()
|
| 1006 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1007 |
+
|
| 1008 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1009 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1010 |
+
|
| 1011 |
+
# Prepare inputs by adding necessary tokens
|
| 1012 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1013 |
+
|
| 1014 |
+
# Update labels tensor for action mask computation later
|
| 1015 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1016 |
+
|
| 1017 |
+
# Get input embeddings and action masks
|
| 1018 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1019 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1020 |
+
|
| 1021 |
+
# Extract language embeddings
|
| 1022 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1023 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1024 |
+
)
|
| 1025 |
+
|
| 1026 |
+
# Process vision features
|
| 1027 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1028 |
+
|
| 1029 |
+
# Add proprioceptive features if provided
|
| 1030 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1031 |
+
if use_proprio:
|
| 1032 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1033 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1034 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1035 |
+
)
|
| 1036 |
+
|
| 1037 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1038 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1039 |
+
|
| 1040 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1041 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1042 |
+
if use_proprio:
|
| 1043 |
+
NUM_PATCHES += 1
|
| 1044 |
+
|
| 1045 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1046 |
+
input_embeddings,
|
| 1047 |
+
all_actions_mask,
|
| 1048 |
+
projected_patch_embeddings,
|
| 1049 |
+
attention_mask,
|
| 1050 |
+
labels,
|
| 1051 |
+
NUM_PATCHES,
|
| 1052 |
+
NUM_PROMPT_TOKENS,
|
| 1053 |
+
action_head,
|
| 1054 |
+
)
|
| 1055 |
+
|
| 1056 |
+
# Unnormalize predicted actions
|
| 1057 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1058 |
+
|
| 1059 |
+
return actions, actions_hidden_states
|
| 1060 |
+
|
| 1061 |
+
@staticmethod
|
| 1062 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1063 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1064 |
+
if unnorm_key is None:
|
| 1065 |
+
assert len(norm_stats) == 1, (
|
| 1066 |
+
f"Your model was trained on more than one dataset, "
|
| 1067 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1068 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1069 |
+
)
|
| 1070 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1071 |
+
# norm states没有加载libero, 为什么?
|
| 1072 |
+
assert unnorm_key in norm_stats, (
|
| 1073 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1074 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1075 |
+
)
|
| 1076 |
+
return unnorm_key
|
| 1077 |
+
|
| 1078 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1079 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1080 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1081 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1082 |
+
|
| 1083 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1084 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1085 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1086 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1087 |
+
|
| 1088 |
+
|
| 1089 |
+
def lisa_forward(
|
| 1090 |
+
self,
|
| 1091 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1092 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1093 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1094 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1095 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1096 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1097 |
+
use_cache: Optional[bool] = None,
|
| 1098 |
+
output_attentions: Optional[bool] = None,
|
| 1099 |
+
output_hidden_states: Optional[bool] = None,
|
| 1100 |
+
output_projector_features: Optional[bool] = None,
|
| 1101 |
+
return_dict: Optional[bool] = None,
|
| 1102 |
+
proprio=None,
|
| 1103 |
+
proprio_projector=None,
|
| 1104 |
+
**kwargs
|
| 1105 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1106 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1107 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1108 |
+
output_hidden_states = (
|
| 1109 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1110 |
+
)
|
| 1111 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1112 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1113 |
+
|
| 1114 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1115 |
+
use_cache = use_cache and not self.training
|
| 1116 |
+
|
| 1117 |
+
# Instantiate Placeholder for Projector Features
|
| 1118 |
+
projected_patch_embeddings = None
|
| 1119 |
+
|
| 1120 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1121 |
+
if input_ids.shape[1] == 1:
|
| 1122 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1123 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1124 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1125 |
+
|
| 1126 |
+
language_model_output = self.language_model(
|
| 1127 |
+
input_ids=input_ids,
|
| 1128 |
+
attention_mask=None,
|
| 1129 |
+
position_ids=None,
|
| 1130 |
+
past_key_values=past_key_values,
|
| 1131 |
+
inputs_embeds=None,
|
| 1132 |
+
labels=None,
|
| 1133 |
+
use_cache=use_cache,
|
| 1134 |
+
output_attentions=output_attentions,
|
| 1135 |
+
output_hidden_states=output_hidden_states,
|
| 1136 |
+
return_dict=return_dict,
|
| 1137 |
+
)
|
| 1138 |
+
|
| 1139 |
+
# === Handle Unimodal Forward ===
|
| 1140 |
+
elif pixel_values is None:
|
| 1141 |
+
raise NotImplementedError
|
| 1142 |
+
|
| 1143 |
+
# === Handle Multimodal Forward ===
|
| 1144 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1145 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1146 |
+
|
| 1147 |
+
# Get input embeddings (from language model embeddings)
|
| 1148 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1149 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1150 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1151 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1152 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1153 |
+
# Get visual features
|
| 1154 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1155 |
+
|
| 1156 |
+
# Add proprioceptive state if provided
|
| 1157 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1158 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1159 |
+
)
|
| 1160 |
+
|
| 1161 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1162 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1163 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1164 |
+
|
| 1165 |
+
# Build multimodal embeddings & attention mask
|
| 1166 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1167 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1168 |
+
)
|
| 1169 |
+
|
| 1170 |
+
# Build labels for multimodal sequence if needed
|
| 1171 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1172 |
+
|
| 1173 |
+
# Dispatch to language model
|
| 1174 |
+
language_model_output = self.language_model(
|
| 1175 |
+
input_ids=None,
|
| 1176 |
+
attention_mask=multimodal_attention_mask,
|
| 1177 |
+
position_ids=None,
|
| 1178 |
+
past_key_values=None,
|
| 1179 |
+
inputs_embeds=multimodal_embeddings,
|
| 1180 |
+
labels=multimodal_labels,
|
| 1181 |
+
use_cache=use_cache,
|
| 1182 |
+
output_attentions=output_attentions,
|
| 1183 |
+
output_hidden_states=output_hidden_states,
|
| 1184 |
+
return_dict=return_dict,
|
| 1185 |
+
)
|
| 1186 |
+
|
| 1187 |
+
|
| 1188 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1189 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1190 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1191 |
+
|
| 1192 |
+
else:
|
| 1193 |
+
raise ValueError(
|
| 1194 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1195 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1196 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1197 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1198 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1199 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1200 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1201 |
+
f"=> `use_cache` = {use_cache}"
|
| 1202 |
+
)
|
| 1203 |
+
|
| 1204 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1205 |
+
if not return_dict:
|
| 1206 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1207 |
+
return *language_model_output, projected_patch_embeddings
|
| 1208 |
+
|
| 1209 |
+
return language_model_output
|
| 1210 |
+
|
| 1211 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1212 |
+
loss=language_model_output.loss,
|
| 1213 |
+
logits=language_model_output.logits,
|
| 1214 |
+
past_key_values=language_model_output.past_key_values,
|
| 1215 |
+
hidden_states=language_model_output.hidden_states,
|
| 1216 |
+
attentions=language_model_output.attentions,
|
| 1217 |
+
projector_features=projected_patch_embeddings,
|
| 1218 |
+
)
|
| 1219 |
+
|
| 1220 |
+
def lisa_predict_action(
|
| 1221 |
+
self,
|
| 1222 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1223 |
+
unnorm_key: Optional[str] = None,
|
| 1224 |
+
proprio=None,
|
| 1225 |
+
proprio_projector=None,
|
| 1226 |
+
action_head:L1RegressionActionHead=None,
|
| 1227 |
+
noisy_action_projector=None,
|
| 1228 |
+
use_film: bool = False,
|
| 1229 |
+
**kwargs: str,
|
| 1230 |
+
) -> np.ndarray:
|
| 1231 |
+
|
| 1232 |
+
pixel_values = kwargs["pixel_values"]
|
| 1233 |
+
attention_mask = kwargs["attention_mask"]
|
| 1234 |
+
|
| 1235 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1236 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1237 |
+
|
| 1238 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1239 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1240 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1241 |
+
if cfg.preset:
|
| 1242 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1243 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1244 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1245 |
+
self.preset = True
|
| 1246 |
+
else:
|
| 1247 |
+
self.preset = False
|
| 1248 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1249 |
+
|
| 1250 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1251 |
+
|
| 1252 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1253 |
+
if use_proprio:
|
| 1254 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1255 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1256 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1257 |
+
)
|
| 1258 |
+
|
| 1259 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1260 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1261 |
+
# if use_proprio:
|
| 1262 |
+
# NUM_PATCHES += 1
|
| 1263 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1264 |
+
input_embeddings,
|
| 1265 |
+
None,
|
| 1266 |
+
projected_patch_embeddings,
|
| 1267 |
+
attention_mask,
|
| 1268 |
+
None, #推理不需要labels
|
| 1269 |
+
None, #推理不需要NUM_PATCHES
|
| 1270 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1271 |
+
action_head,
|
| 1272 |
+
)
|
| 1273 |
+
|
| 1274 |
+
# Unnormalize predicted actions
|
| 1275 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1276 |
+
|
| 1277 |
+
return actions, actions_hidden_states
|
| 1278 |
+
|
| 1279 |
+
|
| 1280 |
+
def mul_predict_action(
|
| 1281 |
+
self,
|
| 1282 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1283 |
+
unnorm_key: Optional[str] = None,
|
| 1284 |
+
proprio=None,
|
| 1285 |
+
proprio_projector=None,
|
| 1286 |
+
action_head:L1RegressionActionHead=None,
|
| 1287 |
+
noisy_action_projector=None,
|
| 1288 |
+
use_film: bool = False,
|
| 1289 |
+
**kwargs: str,
|
| 1290 |
+
) -> np.ndarray:
|
| 1291 |
+
|
| 1292 |
+
pixel_values = kwargs["pixel_values"]
|
| 1293 |
+
attention_mask = kwargs["attention_mask"]
|
| 1294 |
+
|
| 1295 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1296 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1297 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1298 |
+
|
| 1299 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1300 |
+
|
| 1301 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1302 |
+
if use_proprio:
|
| 1303 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1304 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1305 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1306 |
+
)
|
| 1307 |
+
|
| 1308 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1309 |
+
input_embeddings,
|
| 1310 |
+
None,
|
| 1311 |
+
projected_patch_embeddings,
|
| 1312 |
+
attention_mask,
|
| 1313 |
+
None, #推理不需要labels
|
| 1314 |
+
None, #推理不需要NUM_PATCHES
|
| 1315 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1316 |
+
action_head,
|
| 1317 |
+
)
|
| 1318 |
+
|
| 1319 |
+
# Unnormalize predicted actions
|
| 1320 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1321 |
+
|
| 1322 |
+
return actions, actions_hidden_states
|
| 1323 |
+
|
modeling_prismatic.py.back.20250403_122246
ADDED
|
@@ -0,0 +1,1323 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
# 0 , is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 957 |
+
actions_hidden_states = language_model_output.hidden_states[1][-1] # 第一个是29871, 应该acc低, 第二个act, 应该acc高.
|
| 958 |
+
# import ipdb; ipdb.set_trace()
|
| 959 |
+
# actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 960 |
+
actions_hidden_states = actions_hidden_states[:, -1]
|
| 961 |
+
|
| 962 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 963 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 964 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 965 |
+
|
| 966 |
+
return normalized_actions, actions_hidden_states
|
| 967 |
+
|
| 968 |
+
def predict_action(
|
| 969 |
+
self,
|
| 970 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 971 |
+
unnorm_key: Optional[str] = None,
|
| 972 |
+
proprio=None,
|
| 973 |
+
proprio_projector=None,
|
| 974 |
+
action_head=None,
|
| 975 |
+
noisy_action_projector=None,
|
| 976 |
+
use_film: bool = False,
|
| 977 |
+
**kwargs: str,
|
| 978 |
+
) -> np.ndarray:
|
| 979 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 980 |
+
|
| 981 |
+
Args:
|
| 982 |
+
input_ids: Input token ids
|
| 983 |
+
unnorm_key: Key for unnormalization statistics
|
| 984 |
+
proprio: Proprioceptive features
|
| 985 |
+
proprio_projector: Projector for proprioceptive features
|
| 986 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 987 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 988 |
+
use_film: Whether to use FiLM conditioning
|
| 989 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 990 |
+
|
| 991 |
+
Returns:
|
| 992 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 993 |
+
"""
|
| 994 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 995 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 996 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 997 |
+
input_ids = torch.cat(
|
| 998 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 999 |
+
)
|
| 1000 |
+
|
| 1001 |
+
pixel_values = kwargs["pixel_values"]
|
| 1002 |
+
attention_mask = kwargs["attention_mask"]
|
| 1003 |
+
|
| 1004 |
+
# Create fake labels tensor (needed for action mask)
|
| 1005 |
+
labels = input_ids.clone()
|
| 1006 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1007 |
+
|
| 1008 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1009 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1010 |
+
|
| 1011 |
+
# Prepare inputs by adding necessary tokens
|
| 1012 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1013 |
+
|
| 1014 |
+
# Update labels tensor for action mask computation later
|
| 1015 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1016 |
+
|
| 1017 |
+
# Get input embeddings and action masks
|
| 1018 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1019 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1020 |
+
|
| 1021 |
+
# Extract language embeddings
|
| 1022 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1023 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1024 |
+
)
|
| 1025 |
+
|
| 1026 |
+
# Process vision features
|
| 1027 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1028 |
+
|
| 1029 |
+
# Add proprioceptive features if provided
|
| 1030 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1031 |
+
if use_proprio:
|
| 1032 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1033 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1034 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1035 |
+
)
|
| 1036 |
+
|
| 1037 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1038 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1039 |
+
|
| 1040 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1041 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1042 |
+
if use_proprio:
|
| 1043 |
+
NUM_PATCHES += 1
|
| 1044 |
+
|
| 1045 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1046 |
+
input_embeddings,
|
| 1047 |
+
all_actions_mask,
|
| 1048 |
+
projected_patch_embeddings,
|
| 1049 |
+
attention_mask,
|
| 1050 |
+
labels,
|
| 1051 |
+
NUM_PATCHES,
|
| 1052 |
+
NUM_PROMPT_TOKENS,
|
| 1053 |
+
action_head,
|
| 1054 |
+
)
|
| 1055 |
+
|
| 1056 |
+
# Unnormalize predicted actions
|
| 1057 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1058 |
+
|
| 1059 |
+
return actions, actions_hidden_states
|
| 1060 |
+
|
| 1061 |
+
@staticmethod
|
| 1062 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1063 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1064 |
+
if unnorm_key is None:
|
| 1065 |
+
assert len(norm_stats) == 1, (
|
| 1066 |
+
f"Your model was trained on more than one dataset, "
|
| 1067 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1068 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1069 |
+
)
|
| 1070 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1071 |
+
# norm states没有加载libero, 为什么?
|
| 1072 |
+
assert unnorm_key in norm_stats, (
|
| 1073 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1074 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1075 |
+
)
|
| 1076 |
+
return unnorm_key
|
| 1077 |
+
|
| 1078 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1079 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1080 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1081 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1082 |
+
|
| 1083 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1084 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1085 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1086 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1087 |
+
|
| 1088 |
+
|
| 1089 |
+
def lisa_forward(
|
| 1090 |
+
self,
|
| 1091 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1092 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1093 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1094 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1095 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1096 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1097 |
+
use_cache: Optional[bool] = None,
|
| 1098 |
+
output_attentions: Optional[bool] = None,
|
| 1099 |
+
output_hidden_states: Optional[bool] = None,
|
| 1100 |
+
output_projector_features: Optional[bool] = None,
|
| 1101 |
+
return_dict: Optional[bool] = None,
|
| 1102 |
+
proprio=None,
|
| 1103 |
+
proprio_projector=None,
|
| 1104 |
+
**kwargs
|
| 1105 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1106 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1107 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1108 |
+
output_hidden_states = (
|
| 1109 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1110 |
+
)
|
| 1111 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1112 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1113 |
+
|
| 1114 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1115 |
+
use_cache = use_cache and not self.training
|
| 1116 |
+
|
| 1117 |
+
# Instantiate Placeholder for Projector Features
|
| 1118 |
+
projected_patch_embeddings = None
|
| 1119 |
+
|
| 1120 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1121 |
+
if input_ids.shape[1] == 1:
|
| 1122 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1123 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1124 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1125 |
+
|
| 1126 |
+
language_model_output = self.language_model(
|
| 1127 |
+
input_ids=input_ids,
|
| 1128 |
+
attention_mask=None,
|
| 1129 |
+
position_ids=None,
|
| 1130 |
+
past_key_values=past_key_values,
|
| 1131 |
+
inputs_embeds=None,
|
| 1132 |
+
labels=None,
|
| 1133 |
+
use_cache=use_cache,
|
| 1134 |
+
output_attentions=output_attentions,
|
| 1135 |
+
output_hidden_states=output_hidden_states,
|
| 1136 |
+
return_dict=return_dict,
|
| 1137 |
+
)
|
| 1138 |
+
|
| 1139 |
+
# === Handle Unimodal Forward ===
|
| 1140 |
+
elif pixel_values is None:
|
| 1141 |
+
raise NotImplementedError
|
| 1142 |
+
|
| 1143 |
+
# === Handle Multimodal Forward ===
|
| 1144 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1145 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1146 |
+
|
| 1147 |
+
# Get input embeddings (from language model embeddings)
|
| 1148 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1149 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1150 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1151 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1152 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1153 |
+
# Get visual features
|
| 1154 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1155 |
+
|
| 1156 |
+
# Add proprioceptive state if provided
|
| 1157 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1158 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1159 |
+
)
|
| 1160 |
+
|
| 1161 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1162 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1163 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1164 |
+
|
| 1165 |
+
# Build multimodal embeddings & attention mask
|
| 1166 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1167 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1168 |
+
)
|
| 1169 |
+
|
| 1170 |
+
# Build labels for multimodal sequence if needed
|
| 1171 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1172 |
+
|
| 1173 |
+
# Dispatch to language model
|
| 1174 |
+
language_model_output = self.language_model(
|
| 1175 |
+
input_ids=None,
|
| 1176 |
+
attention_mask=multimodal_attention_mask,
|
| 1177 |
+
position_ids=None,
|
| 1178 |
+
past_key_values=None,
|
| 1179 |
+
inputs_embeds=multimodal_embeddings,
|
| 1180 |
+
labels=multimodal_labels,
|
| 1181 |
+
use_cache=use_cache,
|
| 1182 |
+
output_attentions=output_attentions,
|
| 1183 |
+
output_hidden_states=output_hidden_states,
|
| 1184 |
+
return_dict=return_dict,
|
| 1185 |
+
)
|
| 1186 |
+
|
| 1187 |
+
|
| 1188 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1189 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1190 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1191 |
+
|
| 1192 |
+
else:
|
| 1193 |
+
raise ValueError(
|
| 1194 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1195 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1196 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1197 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1198 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1199 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1200 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1201 |
+
f"=> `use_cache` = {use_cache}"
|
| 1202 |
+
)
|
| 1203 |
+
|
| 1204 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1205 |
+
if not return_dict:
|
| 1206 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1207 |
+
return *language_model_output, projected_patch_embeddings
|
| 1208 |
+
|
| 1209 |
+
return language_model_output
|
| 1210 |
+
|
| 1211 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1212 |
+
loss=language_model_output.loss,
|
| 1213 |
+
logits=language_model_output.logits,
|
| 1214 |
+
past_key_values=language_model_output.past_key_values,
|
| 1215 |
+
hidden_states=language_model_output.hidden_states,
|
| 1216 |
+
attentions=language_model_output.attentions,
|
| 1217 |
+
projector_features=projected_patch_embeddings,
|
| 1218 |
+
)
|
| 1219 |
+
|
| 1220 |
+
def lisa_predict_action(
|
| 1221 |
+
self,
|
| 1222 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1223 |
+
unnorm_key: Optional[str] = None,
|
| 1224 |
+
proprio=None,
|
| 1225 |
+
proprio_projector=None,
|
| 1226 |
+
action_head:L1RegressionActionHead=None,
|
| 1227 |
+
noisy_action_projector=None,
|
| 1228 |
+
use_film: bool = False,
|
| 1229 |
+
**kwargs: str,
|
| 1230 |
+
) -> np.ndarray:
|
| 1231 |
+
|
| 1232 |
+
pixel_values = kwargs["pixel_values"]
|
| 1233 |
+
attention_mask = kwargs["attention_mask"]
|
| 1234 |
+
|
| 1235 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1236 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1237 |
+
|
| 1238 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1239 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1240 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1241 |
+
if cfg.preset:
|
| 1242 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1243 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1244 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1245 |
+
self.preset = True
|
| 1246 |
+
else:
|
| 1247 |
+
self.preset = False
|
| 1248 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1249 |
+
|
| 1250 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1251 |
+
|
| 1252 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1253 |
+
if use_proprio:
|
| 1254 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1255 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1256 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1257 |
+
)
|
| 1258 |
+
|
| 1259 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1260 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1261 |
+
# if use_proprio:
|
| 1262 |
+
# NUM_PATCHES += 1
|
| 1263 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1264 |
+
input_embeddings,
|
| 1265 |
+
None,
|
| 1266 |
+
projected_patch_embeddings,
|
| 1267 |
+
attention_mask,
|
| 1268 |
+
None, #推理不需要labels
|
| 1269 |
+
None, #推理不需要NUM_PATCHES
|
| 1270 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1271 |
+
action_head,
|
| 1272 |
+
)
|
| 1273 |
+
|
| 1274 |
+
# Unnormalize predicted actions
|
| 1275 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1276 |
+
|
| 1277 |
+
return actions, actions_hidden_states
|
| 1278 |
+
|
| 1279 |
+
|
| 1280 |
+
def mul_predict_action(
|
| 1281 |
+
self,
|
| 1282 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1283 |
+
unnorm_key: Optional[str] = None,
|
| 1284 |
+
proprio=None,
|
| 1285 |
+
proprio_projector=None,
|
| 1286 |
+
action_head:L1RegressionActionHead=None,
|
| 1287 |
+
noisy_action_projector=None,
|
| 1288 |
+
use_film: bool = False,
|
| 1289 |
+
**kwargs: str,
|
| 1290 |
+
) -> np.ndarray:
|
| 1291 |
+
|
| 1292 |
+
pixel_values = kwargs["pixel_values"]
|
| 1293 |
+
attention_mask = kwargs["attention_mask"]
|
| 1294 |
+
|
| 1295 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1296 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1297 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1298 |
+
|
| 1299 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1300 |
+
|
| 1301 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1302 |
+
if use_proprio:
|
| 1303 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1304 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1305 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1306 |
+
)
|
| 1307 |
+
|
| 1308 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1309 |
+
input_embeddings,
|
| 1310 |
+
None,
|
| 1311 |
+
projected_patch_embeddings,
|
| 1312 |
+
attention_mask,
|
| 1313 |
+
None, #推理不需要labels
|
| 1314 |
+
None, #推理不需要NUM_PATCHES
|
| 1315 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1316 |
+
action_head,
|
| 1317 |
+
)
|
| 1318 |
+
|
| 1319 |
+
# Unnormalize predicted actions
|
| 1320 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1321 |
+
|
| 1322 |
+
return actions, actions_hidden_states
|
| 1323 |
+
|
modeling_prismatic.py.back.20250403_125446
ADDED
|
@@ -0,0 +1,1321 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
# 0 , is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 957 |
+
actions_hidden_states = language_model_output.hidden_states[1][-1]
|
| 958 |
+
actions_hidden_states = actions_hidden_states[:, -1]
|
| 959 |
+
|
| 960 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 961 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 962 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 963 |
+
|
| 964 |
+
return normalized_actions, actions_hidden_states
|
| 965 |
+
|
| 966 |
+
def predict_action(
|
| 967 |
+
self,
|
| 968 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 969 |
+
unnorm_key: Optional[str] = None,
|
| 970 |
+
proprio=None,
|
| 971 |
+
proprio_projector=None,
|
| 972 |
+
action_head=None,
|
| 973 |
+
noisy_action_projector=None,
|
| 974 |
+
use_film: bool = False,
|
| 975 |
+
**kwargs: str,
|
| 976 |
+
) -> np.ndarray:
|
| 977 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 978 |
+
|
| 979 |
+
Args:
|
| 980 |
+
input_ids: Input token ids
|
| 981 |
+
unnorm_key: Key for unnormalization statistics
|
| 982 |
+
proprio: Proprioceptive features
|
| 983 |
+
proprio_projector: Projector for proprioceptive features
|
| 984 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 985 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 986 |
+
use_film: Whether to use FiLM conditioning
|
| 987 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 988 |
+
|
| 989 |
+
Returns:
|
| 990 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 991 |
+
"""
|
| 992 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 993 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 994 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 995 |
+
input_ids = torch.cat(
|
| 996 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 997 |
+
)
|
| 998 |
+
|
| 999 |
+
pixel_values = kwargs["pixel_values"]
|
| 1000 |
+
attention_mask = kwargs["attention_mask"]
|
| 1001 |
+
|
| 1002 |
+
# Create fake labels tensor (needed for action mask)
|
| 1003 |
+
labels = input_ids.clone()
|
| 1004 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1005 |
+
|
| 1006 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1007 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1008 |
+
|
| 1009 |
+
# Prepare inputs by adding necessary tokens
|
| 1010 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1011 |
+
|
| 1012 |
+
# Update labels tensor for action mask computation later
|
| 1013 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1014 |
+
|
| 1015 |
+
# Get input embeddings and action masks
|
| 1016 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1017 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1018 |
+
|
| 1019 |
+
# Extract language embeddings
|
| 1020 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1021 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1022 |
+
)
|
| 1023 |
+
|
| 1024 |
+
# Process vision features
|
| 1025 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1026 |
+
|
| 1027 |
+
# Add proprioceptive features if provided
|
| 1028 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1029 |
+
if use_proprio:
|
| 1030 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1031 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1032 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1033 |
+
)
|
| 1034 |
+
|
| 1035 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1036 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1037 |
+
|
| 1038 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1039 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1040 |
+
if use_proprio:
|
| 1041 |
+
NUM_PATCHES += 1
|
| 1042 |
+
|
| 1043 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1044 |
+
input_embeddings,
|
| 1045 |
+
all_actions_mask,
|
| 1046 |
+
projected_patch_embeddings,
|
| 1047 |
+
attention_mask,
|
| 1048 |
+
labels,
|
| 1049 |
+
NUM_PATCHES,
|
| 1050 |
+
NUM_PROMPT_TOKENS,
|
| 1051 |
+
action_head,
|
| 1052 |
+
)
|
| 1053 |
+
|
| 1054 |
+
# Unnormalize predicted actions
|
| 1055 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1056 |
+
|
| 1057 |
+
return actions, actions_hidden_states
|
| 1058 |
+
|
| 1059 |
+
@staticmethod
|
| 1060 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1061 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1062 |
+
if unnorm_key is None:
|
| 1063 |
+
assert len(norm_stats) == 1, (
|
| 1064 |
+
f"Your model was trained on more than one dataset, "
|
| 1065 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1066 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1067 |
+
)
|
| 1068 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1069 |
+
# norm states没有加载libero, 为什么?
|
| 1070 |
+
assert unnorm_key in norm_stats, (
|
| 1071 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1072 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1073 |
+
)
|
| 1074 |
+
return unnorm_key
|
| 1075 |
+
|
| 1076 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1077 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1078 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1079 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1080 |
+
|
| 1081 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1082 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1083 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1084 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1085 |
+
|
| 1086 |
+
|
| 1087 |
+
def lisa_forward(
|
| 1088 |
+
self,
|
| 1089 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1090 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1091 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1092 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1093 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1094 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1095 |
+
use_cache: Optional[bool] = None,
|
| 1096 |
+
output_attentions: Optional[bool] = None,
|
| 1097 |
+
output_hidden_states: Optional[bool] = None,
|
| 1098 |
+
output_projector_features: Optional[bool] = None,
|
| 1099 |
+
return_dict: Optional[bool] = None,
|
| 1100 |
+
proprio=None,
|
| 1101 |
+
proprio_projector=None,
|
| 1102 |
+
**kwargs
|
| 1103 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1104 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1105 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1106 |
+
output_hidden_states = (
|
| 1107 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1108 |
+
)
|
| 1109 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1110 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1111 |
+
|
| 1112 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1113 |
+
use_cache = use_cache and not self.training
|
| 1114 |
+
|
| 1115 |
+
# Instantiate Placeholder for Projector Features
|
| 1116 |
+
projected_patch_embeddings = None
|
| 1117 |
+
|
| 1118 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1119 |
+
if input_ids.shape[1] == 1:
|
| 1120 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1121 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1122 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1123 |
+
|
| 1124 |
+
language_model_output = self.language_model(
|
| 1125 |
+
input_ids=input_ids,
|
| 1126 |
+
attention_mask=None,
|
| 1127 |
+
position_ids=None,
|
| 1128 |
+
past_key_values=past_key_values,
|
| 1129 |
+
inputs_embeds=None,
|
| 1130 |
+
labels=None,
|
| 1131 |
+
use_cache=use_cache,
|
| 1132 |
+
output_attentions=output_attentions,
|
| 1133 |
+
output_hidden_states=output_hidden_states,
|
| 1134 |
+
return_dict=return_dict,
|
| 1135 |
+
)
|
| 1136 |
+
|
| 1137 |
+
# === Handle Unimodal Forward ===
|
| 1138 |
+
elif pixel_values is None:
|
| 1139 |
+
raise NotImplementedError
|
| 1140 |
+
|
| 1141 |
+
# === Handle Multimodal Forward ===
|
| 1142 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1143 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1144 |
+
|
| 1145 |
+
# Get input embeddings (from language model embeddings)
|
| 1146 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1147 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1148 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1149 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1150 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1151 |
+
# Get visual features
|
| 1152 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1153 |
+
|
| 1154 |
+
# Add proprioceptive state if provided
|
| 1155 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1156 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1157 |
+
)
|
| 1158 |
+
|
| 1159 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1160 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1161 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1162 |
+
|
| 1163 |
+
# Build multimodal embeddings & attention mask
|
| 1164 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1165 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1166 |
+
)
|
| 1167 |
+
|
| 1168 |
+
# Build labels for multimodal sequence if needed
|
| 1169 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1170 |
+
|
| 1171 |
+
# Dispatch to language model
|
| 1172 |
+
language_model_output = self.language_model(
|
| 1173 |
+
input_ids=None,
|
| 1174 |
+
attention_mask=multimodal_attention_mask,
|
| 1175 |
+
position_ids=None,
|
| 1176 |
+
past_key_values=None,
|
| 1177 |
+
inputs_embeds=multimodal_embeddings,
|
| 1178 |
+
labels=multimodal_labels,
|
| 1179 |
+
use_cache=use_cache,
|
| 1180 |
+
output_attentions=output_attentions,
|
| 1181 |
+
output_hidden_states=output_hidden_states,
|
| 1182 |
+
return_dict=return_dict,
|
| 1183 |
+
)
|
| 1184 |
+
|
| 1185 |
+
|
| 1186 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1187 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1188 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1189 |
+
|
| 1190 |
+
else:
|
| 1191 |
+
raise ValueError(
|
| 1192 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1193 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1194 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1195 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1196 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1197 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1198 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1199 |
+
f"=> `use_cache` = {use_cache}"
|
| 1200 |
+
)
|
| 1201 |
+
|
| 1202 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1203 |
+
if not return_dict:
|
| 1204 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1205 |
+
return *language_model_output, projected_patch_embeddings
|
| 1206 |
+
|
| 1207 |
+
return language_model_output
|
| 1208 |
+
|
| 1209 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1210 |
+
loss=language_model_output.loss,
|
| 1211 |
+
logits=language_model_output.logits,
|
| 1212 |
+
past_key_values=language_model_output.past_key_values,
|
| 1213 |
+
hidden_states=language_model_output.hidden_states,
|
| 1214 |
+
attentions=language_model_output.attentions,
|
| 1215 |
+
projector_features=projected_patch_embeddings,
|
| 1216 |
+
)
|
| 1217 |
+
|
| 1218 |
+
def lisa_predict_action(
|
| 1219 |
+
self,
|
| 1220 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1221 |
+
unnorm_key: Optional[str] = None,
|
| 1222 |
+
proprio=None,
|
| 1223 |
+
proprio_projector=None,
|
| 1224 |
+
action_head:L1RegressionActionHead=None,
|
| 1225 |
+
noisy_action_projector=None,
|
| 1226 |
+
use_film: bool = False,
|
| 1227 |
+
**kwargs: str,
|
| 1228 |
+
) -> np.ndarray:
|
| 1229 |
+
|
| 1230 |
+
pixel_values = kwargs["pixel_values"]
|
| 1231 |
+
attention_mask = kwargs["attention_mask"]
|
| 1232 |
+
|
| 1233 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1234 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1235 |
+
|
| 1236 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1237 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1238 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1239 |
+
if cfg.preset:
|
| 1240 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1241 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1242 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1243 |
+
self.preset = True
|
| 1244 |
+
else:
|
| 1245 |
+
self.preset = False
|
| 1246 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1247 |
+
|
| 1248 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1249 |
+
|
| 1250 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1251 |
+
if use_proprio:
|
| 1252 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1253 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1254 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1255 |
+
)
|
| 1256 |
+
|
| 1257 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1258 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1259 |
+
# if use_proprio:
|
| 1260 |
+
# NUM_PATCHES += 1
|
| 1261 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1262 |
+
input_embeddings,
|
| 1263 |
+
None,
|
| 1264 |
+
projected_patch_embeddings,
|
| 1265 |
+
attention_mask,
|
| 1266 |
+
None, #推理不需要labels
|
| 1267 |
+
None, #推理不需要NUM_PATCHES
|
| 1268 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1269 |
+
action_head,
|
| 1270 |
+
)
|
| 1271 |
+
|
| 1272 |
+
# Unnormalize predicted actions
|
| 1273 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1274 |
+
|
| 1275 |
+
return actions, actions_hidden_states
|
| 1276 |
+
|
| 1277 |
+
|
| 1278 |
+
def mul_predict_action(
|
| 1279 |
+
self,
|
| 1280 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1281 |
+
unnorm_key: Optional[str] = None,
|
| 1282 |
+
proprio=None,
|
| 1283 |
+
proprio_projector=None,
|
| 1284 |
+
action_head:L1RegressionActionHead=None,
|
| 1285 |
+
noisy_action_projector=None,
|
| 1286 |
+
use_film: bool = False,
|
| 1287 |
+
**kwargs: str,
|
| 1288 |
+
) -> np.ndarray:
|
| 1289 |
+
|
| 1290 |
+
pixel_values = kwargs["pixel_values"]
|
| 1291 |
+
attention_mask = kwargs["attention_mask"]
|
| 1292 |
+
|
| 1293 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1294 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1295 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1296 |
+
|
| 1297 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1298 |
+
|
| 1299 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1300 |
+
if use_proprio:
|
| 1301 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1302 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1303 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1304 |
+
)
|
| 1305 |
+
|
| 1306 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1307 |
+
input_embeddings,
|
| 1308 |
+
None,
|
| 1309 |
+
projected_patch_embeddings,
|
| 1310 |
+
attention_mask,
|
| 1311 |
+
None, #推理不需要labels
|
| 1312 |
+
None, #推理不需要NUM_PATCHES
|
| 1313 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1314 |
+
action_head,
|
| 1315 |
+
)
|
| 1316 |
+
|
| 1317 |
+
# Unnormalize predicted actions
|
| 1318 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1319 |
+
|
| 1320 |
+
return actions, actions_hidden_states
|
| 1321 |
+
|
modeling_prismatic.py.back.20250403_125952
ADDED
|
@@ -0,0 +1,1327 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
# 0 , is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 957 |
+
# actions_hidden_states = language_model_output.hidden_states[1][-1]
|
| 958 |
+
actions_hidden_states = language_model_output.hidden_states[0][-1]
|
| 959 |
+
# actions_hidden_states = actions_hidden_states[:, -1]
|
| 960 |
+
|
| 961 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 962 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 963 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 964 |
+
|
| 965 |
+
return normalized_actions, actions_hidden_states
|
| 966 |
+
|
| 967 |
+
def predict_action(
|
| 968 |
+
self,
|
| 969 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 970 |
+
unnorm_key: Optional[str] = None,
|
| 971 |
+
proprio=None,
|
| 972 |
+
proprio_projector=None,
|
| 973 |
+
action_head=None,
|
| 974 |
+
noisy_action_projector=None,
|
| 975 |
+
use_film: bool = False,
|
| 976 |
+
**kwargs: str,
|
| 977 |
+
) -> np.ndarray:
|
| 978 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 979 |
+
|
| 980 |
+
Args:
|
| 981 |
+
input_ids: Input token ids
|
| 982 |
+
unnorm_key: Key for unnormalization statistics
|
| 983 |
+
proprio: Proprioceptive features
|
| 984 |
+
proprio_projector: Projector for proprioceptive features
|
| 985 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 986 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 987 |
+
use_film: Whether to use FiLM conditioning
|
| 988 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 989 |
+
|
| 990 |
+
Returns:
|
| 991 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 992 |
+
"""
|
| 993 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 994 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 995 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 996 |
+
input_ids = torch.cat(
|
| 997 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 998 |
+
)
|
| 999 |
+
|
| 1000 |
+
pixel_values = kwargs["pixel_values"]
|
| 1001 |
+
attention_mask = kwargs["attention_mask"]
|
| 1002 |
+
|
| 1003 |
+
# Create fake labels tensor (needed for action mask)
|
| 1004 |
+
labels = input_ids.clone()
|
| 1005 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1006 |
+
|
| 1007 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1008 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1009 |
+
|
| 1010 |
+
# Prepare inputs by adding necessary tokens
|
| 1011 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1012 |
+
|
| 1013 |
+
# Update labels tensor for action mask computation later
|
| 1014 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1015 |
+
|
| 1016 |
+
# Get input embeddings and action masks
|
| 1017 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1018 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1019 |
+
|
| 1020 |
+
# Extract language embeddings
|
| 1021 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1022 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1023 |
+
)
|
| 1024 |
+
|
| 1025 |
+
# Process vision features
|
| 1026 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1027 |
+
|
| 1028 |
+
# Add proprioceptive features if provided
|
| 1029 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1030 |
+
if use_proprio:
|
| 1031 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1032 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1033 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1034 |
+
)
|
| 1035 |
+
|
| 1036 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1037 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1038 |
+
|
| 1039 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1040 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1041 |
+
if use_proprio:
|
| 1042 |
+
NUM_PATCHES += 1
|
| 1043 |
+
|
| 1044 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1045 |
+
input_embeddings,
|
| 1046 |
+
all_actions_mask,
|
| 1047 |
+
projected_patch_embeddings,
|
| 1048 |
+
attention_mask,
|
| 1049 |
+
labels,
|
| 1050 |
+
NUM_PATCHES,
|
| 1051 |
+
NUM_PROMPT_TOKENS,
|
| 1052 |
+
action_head,
|
| 1053 |
+
)
|
| 1054 |
+
|
| 1055 |
+
# Unnormalize predicted actions
|
| 1056 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1057 |
+
|
| 1058 |
+
return actions, actions_hidden_states
|
| 1059 |
+
|
| 1060 |
+
@staticmethod
|
| 1061 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1062 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1063 |
+
if unnorm_key is None:
|
| 1064 |
+
assert len(norm_stats) == 1, (
|
| 1065 |
+
f"Your model was trained on more than one dataset, "
|
| 1066 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1067 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1068 |
+
)
|
| 1069 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1070 |
+
# norm states没有加载libero, 为什么?
|
| 1071 |
+
assert unnorm_key in norm_stats, (
|
| 1072 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1073 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1074 |
+
)
|
| 1075 |
+
return unnorm_key
|
| 1076 |
+
|
| 1077 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1078 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1079 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1080 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1081 |
+
|
| 1082 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1083 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1084 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1085 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1086 |
+
|
| 1087 |
+
|
| 1088 |
+
def lisa_forward(
|
| 1089 |
+
self,
|
| 1090 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1091 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1092 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1093 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1094 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1095 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1096 |
+
use_cache: Optional[bool] = None,
|
| 1097 |
+
output_attentions: Optional[bool] = None,
|
| 1098 |
+
output_hidden_states: Optional[bool] = None,
|
| 1099 |
+
output_projector_features: Optional[bool] = None,
|
| 1100 |
+
return_dict: Optional[bool] = None,
|
| 1101 |
+
proprio=None,
|
| 1102 |
+
proprio_projector=None,
|
| 1103 |
+
**kwargs
|
| 1104 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1105 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1106 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1107 |
+
output_hidden_states = (
|
| 1108 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1109 |
+
)
|
| 1110 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1111 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1112 |
+
|
| 1113 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1114 |
+
use_cache = use_cache and not self.training
|
| 1115 |
+
|
| 1116 |
+
# Instantiate Placeholder for Projector Features
|
| 1117 |
+
projected_patch_embeddings = None
|
| 1118 |
+
|
| 1119 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1120 |
+
if input_ids.shape[1] == 1:
|
| 1121 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1122 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1123 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1124 |
+
|
| 1125 |
+
language_model_output = self.language_model(
|
| 1126 |
+
input_ids=input_ids,
|
| 1127 |
+
attention_mask=None,
|
| 1128 |
+
position_ids=None,
|
| 1129 |
+
past_key_values=past_key_values,
|
| 1130 |
+
inputs_embeds=None,
|
| 1131 |
+
labels=None,
|
| 1132 |
+
use_cache=use_cache,
|
| 1133 |
+
output_attentions=output_attentions,
|
| 1134 |
+
output_hidden_states=output_hidden_states,
|
| 1135 |
+
return_dict=return_dict,
|
| 1136 |
+
)
|
| 1137 |
+
|
| 1138 |
+
# === Handle Unimodal Forward ===
|
| 1139 |
+
elif pixel_values is None:
|
| 1140 |
+
raise NotImplementedError
|
| 1141 |
+
|
| 1142 |
+
# === Handle Multimodal Forward ===
|
| 1143 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1144 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1145 |
+
|
| 1146 |
+
# Get input embeddings (from language model embeddings)
|
| 1147 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1148 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1149 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1150 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1151 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1152 |
+
# Get visual features
|
| 1153 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1154 |
+
|
| 1155 |
+
# Add proprioceptive state if provided
|
| 1156 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1157 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1158 |
+
)
|
| 1159 |
+
|
| 1160 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1161 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1162 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1163 |
+
|
| 1164 |
+
# Build multimodal embeddings & attention mask
|
| 1165 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1166 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1167 |
+
)
|
| 1168 |
+
|
| 1169 |
+
# Build labels for multimodal sequence if needed
|
| 1170 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1171 |
+
|
| 1172 |
+
# Dispatch to language model
|
| 1173 |
+
language_model_output = self.language_model(
|
| 1174 |
+
input_ids=None,
|
| 1175 |
+
attention_mask=multimodal_attention_mask,
|
| 1176 |
+
position_ids=None,
|
| 1177 |
+
past_key_values=None,
|
| 1178 |
+
inputs_embeds=multimodal_embeddings,
|
| 1179 |
+
labels=multimodal_labels,
|
| 1180 |
+
use_cache=use_cache,
|
| 1181 |
+
output_attentions=output_attentions,
|
| 1182 |
+
output_hidden_states=output_hidden_states,
|
| 1183 |
+
return_dict=return_dict,
|
| 1184 |
+
)
|
| 1185 |
+
|
| 1186 |
+
|
| 1187 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1188 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1189 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1190 |
+
|
| 1191 |
+
else:
|
| 1192 |
+
raise ValueError(
|
| 1193 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1194 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1195 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1196 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1197 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1198 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1199 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1200 |
+
f"=> `use_cache` = {use_cache}"
|
| 1201 |
+
)
|
| 1202 |
+
|
| 1203 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1204 |
+
if not return_dict:
|
| 1205 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1206 |
+
return *language_model_output, projected_patch_embeddings
|
| 1207 |
+
|
| 1208 |
+
return language_model_output
|
| 1209 |
+
|
| 1210 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1211 |
+
loss=language_model_output.loss,
|
| 1212 |
+
logits=language_model_output.logits,
|
| 1213 |
+
past_key_values=language_model_output.past_key_values,
|
| 1214 |
+
hidden_states=language_model_output.hidden_states,
|
| 1215 |
+
attentions=language_model_output.attentions,
|
| 1216 |
+
projector_features=projected_patch_embeddings,
|
| 1217 |
+
)
|
| 1218 |
+
|
| 1219 |
+
def lisa_predict_action(
|
| 1220 |
+
self,
|
| 1221 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1222 |
+
unnorm_key: Optional[str] = None,
|
| 1223 |
+
proprio=None,
|
| 1224 |
+
proprio_projector=None,
|
| 1225 |
+
action_head:L1RegressionActionHead=None,
|
| 1226 |
+
noisy_action_projector=None,
|
| 1227 |
+
use_film: bool = False,
|
| 1228 |
+
**kwargs: str,
|
| 1229 |
+
) -> np.ndarray:
|
| 1230 |
+
|
| 1231 |
+
pixel_values = kwargs["pixel_values"]
|
| 1232 |
+
attention_mask = kwargs["attention_mask"]
|
| 1233 |
+
|
| 1234 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1235 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1236 |
+
|
| 1237 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1238 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1239 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1240 |
+
if cfg.preset:
|
| 1241 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1242 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1243 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1244 |
+
self.preset = True
|
| 1245 |
+
else:
|
| 1246 |
+
self.preset = False
|
| 1247 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1248 |
+
|
| 1249 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1250 |
+
|
| 1251 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1252 |
+
if use_proprio:
|
| 1253 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1254 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1255 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1256 |
+
)
|
| 1257 |
+
|
| 1258 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1259 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1260 |
+
# if use_proprio:
|
| 1261 |
+
# NUM_PATCHES += 1
|
| 1262 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1263 |
+
input_embeddings,
|
| 1264 |
+
None,
|
| 1265 |
+
projected_patch_embeddings,
|
| 1266 |
+
attention_mask,
|
| 1267 |
+
None, #推理不需要labels
|
| 1268 |
+
None, #推理不需要NUM_PATCHES
|
| 1269 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1270 |
+
action_head,
|
| 1271 |
+
)
|
| 1272 |
+
|
| 1273 |
+
# Unnormalize predicted actions
|
| 1274 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1275 |
+
|
| 1276 |
+
return actions, actions_hidden_states
|
| 1277 |
+
|
| 1278 |
+
|
| 1279 |
+
def mul_predict_action(
|
| 1280 |
+
self,
|
| 1281 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1282 |
+
unnorm_key: Optional[str] = None,
|
| 1283 |
+
proprio=None,
|
| 1284 |
+
proprio_projector=None,
|
| 1285 |
+
action_head:L1RegressionActionHead=None,
|
| 1286 |
+
noisy_action_projector=None,
|
| 1287 |
+
use_film: bool = False,
|
| 1288 |
+
**kwargs: str,
|
| 1289 |
+
) -> np.ndarray:
|
| 1290 |
+
|
| 1291 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 1292 |
+
input_ids = torch.cat(
|
| 1293 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 1294 |
+
)
|
| 1295 |
+
|
| 1296 |
+
pixel_values = kwargs["pixel_values"]
|
| 1297 |
+
attention_mask = kwargs["attention_mask"]
|
| 1298 |
+
|
| 1299 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1300 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1301 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1302 |
+
|
| 1303 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1304 |
+
|
| 1305 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1306 |
+
if use_proprio:
|
| 1307 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1308 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1309 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1310 |
+
)
|
| 1311 |
+
|
| 1312 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1313 |
+
input_embeddings,
|
| 1314 |
+
None,
|
| 1315 |
+
projected_patch_embeddings,
|
| 1316 |
+
attention_mask,
|
| 1317 |
+
None, #推理不需要labels
|
| 1318 |
+
None, #推理不需要NUM_PATCHES
|
| 1319 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1320 |
+
action_head,
|
| 1321 |
+
)
|
| 1322 |
+
|
| 1323 |
+
# Unnormalize predicted actions
|
| 1324 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1325 |
+
|
| 1326 |
+
return actions, actions_hidden_states
|
| 1327 |
+
|
modeling_prismatic.py.back.20250403_130155
ADDED
|
@@ -0,0 +1,1328 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
# 0 , is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 957 |
+
# actions_hidden_states = language_model_output.hidden_states[1][-1]
|
| 958 |
+
import ipdb; ipdb.set_trace()
|
| 959 |
+
actions_hidden_states = language_model_output.hidden_states[0][-1]
|
| 960 |
+
# actions_hidden_states = actions_hidden_states[:, -1]
|
| 961 |
+
|
| 962 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 963 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 964 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 965 |
+
|
| 966 |
+
return normalized_actions, actions_hidden_states
|
| 967 |
+
|
| 968 |
+
def predict_action(
|
| 969 |
+
self,
|
| 970 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 971 |
+
unnorm_key: Optional[str] = None,
|
| 972 |
+
proprio=None,
|
| 973 |
+
proprio_projector=None,
|
| 974 |
+
action_head=None,
|
| 975 |
+
noisy_action_projector=None,
|
| 976 |
+
use_film: bool = False,
|
| 977 |
+
**kwargs: str,
|
| 978 |
+
) -> np.ndarray:
|
| 979 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 980 |
+
|
| 981 |
+
Args:
|
| 982 |
+
input_ids: Input token ids
|
| 983 |
+
unnorm_key: Key for unnormalization statistics
|
| 984 |
+
proprio: Proprioceptive features
|
| 985 |
+
proprio_projector: Projector for proprioceptive features
|
| 986 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 987 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 988 |
+
use_film: Whether to use FiLM conditioning
|
| 989 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 990 |
+
|
| 991 |
+
Returns:
|
| 992 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 993 |
+
"""
|
| 994 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 995 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 996 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 997 |
+
input_ids = torch.cat(
|
| 998 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 999 |
+
)
|
| 1000 |
+
|
| 1001 |
+
pixel_values = kwargs["pixel_values"]
|
| 1002 |
+
attention_mask = kwargs["attention_mask"]
|
| 1003 |
+
|
| 1004 |
+
# Create fake labels tensor (needed for action mask)
|
| 1005 |
+
labels = input_ids.clone()
|
| 1006 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1007 |
+
|
| 1008 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1009 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1010 |
+
|
| 1011 |
+
# Prepare inputs by adding necessary tokens
|
| 1012 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1013 |
+
|
| 1014 |
+
# Update labels tensor for action mask computation later
|
| 1015 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1016 |
+
|
| 1017 |
+
# Get input embeddings and action masks
|
| 1018 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1019 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1020 |
+
|
| 1021 |
+
# Extract language embeddings
|
| 1022 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1023 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1024 |
+
)
|
| 1025 |
+
|
| 1026 |
+
# Process vision features
|
| 1027 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1028 |
+
|
| 1029 |
+
# Add proprioceptive features if provided
|
| 1030 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1031 |
+
if use_proprio:
|
| 1032 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1033 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1034 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1035 |
+
)
|
| 1036 |
+
|
| 1037 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1038 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1039 |
+
|
| 1040 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1041 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1042 |
+
if use_proprio:
|
| 1043 |
+
NUM_PATCHES += 1
|
| 1044 |
+
|
| 1045 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1046 |
+
input_embeddings,
|
| 1047 |
+
all_actions_mask,
|
| 1048 |
+
projected_patch_embeddings,
|
| 1049 |
+
attention_mask,
|
| 1050 |
+
labels,
|
| 1051 |
+
NUM_PATCHES,
|
| 1052 |
+
NUM_PROMPT_TOKENS,
|
| 1053 |
+
action_head,
|
| 1054 |
+
)
|
| 1055 |
+
|
| 1056 |
+
# Unnormalize predicted actions
|
| 1057 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1058 |
+
|
| 1059 |
+
return actions, actions_hidden_states
|
| 1060 |
+
|
| 1061 |
+
@staticmethod
|
| 1062 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1063 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1064 |
+
if unnorm_key is None:
|
| 1065 |
+
assert len(norm_stats) == 1, (
|
| 1066 |
+
f"Your model was trained on more than one dataset, "
|
| 1067 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1068 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1069 |
+
)
|
| 1070 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1071 |
+
# norm states没有加载libero, 为什么?
|
| 1072 |
+
assert unnorm_key in norm_stats, (
|
| 1073 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1074 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1075 |
+
)
|
| 1076 |
+
return unnorm_key
|
| 1077 |
+
|
| 1078 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1079 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1080 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1081 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1082 |
+
|
| 1083 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1084 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1085 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1086 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1087 |
+
|
| 1088 |
+
|
| 1089 |
+
def lisa_forward(
|
| 1090 |
+
self,
|
| 1091 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1092 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1093 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1094 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1095 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1096 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1097 |
+
use_cache: Optional[bool] = None,
|
| 1098 |
+
output_attentions: Optional[bool] = None,
|
| 1099 |
+
output_hidden_states: Optional[bool] = None,
|
| 1100 |
+
output_projector_features: Optional[bool] = None,
|
| 1101 |
+
return_dict: Optional[bool] = None,
|
| 1102 |
+
proprio=None,
|
| 1103 |
+
proprio_projector=None,
|
| 1104 |
+
**kwargs
|
| 1105 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1106 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1107 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1108 |
+
output_hidden_states = (
|
| 1109 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1110 |
+
)
|
| 1111 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1112 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1113 |
+
|
| 1114 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1115 |
+
use_cache = use_cache and not self.training
|
| 1116 |
+
|
| 1117 |
+
# Instantiate Placeholder for Projector Features
|
| 1118 |
+
projected_patch_embeddings = None
|
| 1119 |
+
|
| 1120 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1121 |
+
if input_ids.shape[1] == 1:
|
| 1122 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1123 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1124 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1125 |
+
|
| 1126 |
+
language_model_output = self.language_model(
|
| 1127 |
+
input_ids=input_ids,
|
| 1128 |
+
attention_mask=None,
|
| 1129 |
+
position_ids=None,
|
| 1130 |
+
past_key_values=past_key_values,
|
| 1131 |
+
inputs_embeds=None,
|
| 1132 |
+
labels=None,
|
| 1133 |
+
use_cache=use_cache,
|
| 1134 |
+
output_attentions=output_attentions,
|
| 1135 |
+
output_hidden_states=output_hidden_states,
|
| 1136 |
+
return_dict=return_dict,
|
| 1137 |
+
)
|
| 1138 |
+
|
| 1139 |
+
# === Handle Unimodal Forward ===
|
| 1140 |
+
elif pixel_values is None:
|
| 1141 |
+
raise NotImplementedError
|
| 1142 |
+
|
| 1143 |
+
# === Handle Multimodal Forward ===
|
| 1144 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1145 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1146 |
+
|
| 1147 |
+
# Get input embeddings (from language model embeddings)
|
| 1148 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1149 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1150 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1151 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1152 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1153 |
+
# Get visual features
|
| 1154 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1155 |
+
|
| 1156 |
+
# Add proprioceptive state if provided
|
| 1157 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1158 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1159 |
+
)
|
| 1160 |
+
|
| 1161 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1162 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1163 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1164 |
+
|
| 1165 |
+
# Build multimodal embeddings & attention mask
|
| 1166 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1167 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1168 |
+
)
|
| 1169 |
+
|
| 1170 |
+
# Build labels for multimodal sequence if needed
|
| 1171 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1172 |
+
|
| 1173 |
+
# Dispatch to language model
|
| 1174 |
+
language_model_output = self.language_model(
|
| 1175 |
+
input_ids=None,
|
| 1176 |
+
attention_mask=multimodal_attention_mask,
|
| 1177 |
+
position_ids=None,
|
| 1178 |
+
past_key_values=None,
|
| 1179 |
+
inputs_embeds=multimodal_embeddings,
|
| 1180 |
+
labels=multimodal_labels,
|
| 1181 |
+
use_cache=use_cache,
|
| 1182 |
+
output_attentions=output_attentions,
|
| 1183 |
+
output_hidden_states=output_hidden_states,
|
| 1184 |
+
return_dict=return_dict,
|
| 1185 |
+
)
|
| 1186 |
+
|
| 1187 |
+
|
| 1188 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1189 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1190 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1191 |
+
|
| 1192 |
+
else:
|
| 1193 |
+
raise ValueError(
|
| 1194 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1195 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1196 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1197 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1198 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1199 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1200 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1201 |
+
f"=> `use_cache` = {use_cache}"
|
| 1202 |
+
)
|
| 1203 |
+
|
| 1204 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1205 |
+
if not return_dict:
|
| 1206 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1207 |
+
return *language_model_output, projected_patch_embeddings
|
| 1208 |
+
|
| 1209 |
+
return language_model_output
|
| 1210 |
+
|
| 1211 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1212 |
+
loss=language_model_output.loss,
|
| 1213 |
+
logits=language_model_output.logits,
|
| 1214 |
+
past_key_values=language_model_output.past_key_values,
|
| 1215 |
+
hidden_states=language_model_output.hidden_states,
|
| 1216 |
+
attentions=language_model_output.attentions,
|
| 1217 |
+
projector_features=projected_patch_embeddings,
|
| 1218 |
+
)
|
| 1219 |
+
|
| 1220 |
+
def lisa_predict_action(
|
| 1221 |
+
self,
|
| 1222 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1223 |
+
unnorm_key: Optional[str] = None,
|
| 1224 |
+
proprio=None,
|
| 1225 |
+
proprio_projector=None,
|
| 1226 |
+
action_head:L1RegressionActionHead=None,
|
| 1227 |
+
noisy_action_projector=None,
|
| 1228 |
+
use_film: bool = False,
|
| 1229 |
+
**kwargs: str,
|
| 1230 |
+
) -> np.ndarray:
|
| 1231 |
+
|
| 1232 |
+
pixel_values = kwargs["pixel_values"]
|
| 1233 |
+
attention_mask = kwargs["attention_mask"]
|
| 1234 |
+
|
| 1235 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1236 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1237 |
+
|
| 1238 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1239 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1240 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1241 |
+
if cfg.preset:
|
| 1242 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1243 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1244 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1245 |
+
self.preset = True
|
| 1246 |
+
else:
|
| 1247 |
+
self.preset = False
|
| 1248 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1249 |
+
|
| 1250 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1251 |
+
|
| 1252 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1253 |
+
if use_proprio:
|
| 1254 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1255 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1256 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1257 |
+
)
|
| 1258 |
+
|
| 1259 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1260 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1261 |
+
# if use_proprio:
|
| 1262 |
+
# NUM_PATCHES += 1
|
| 1263 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1264 |
+
input_embeddings,
|
| 1265 |
+
None,
|
| 1266 |
+
projected_patch_embeddings,
|
| 1267 |
+
attention_mask,
|
| 1268 |
+
None, #推理不需要labels
|
| 1269 |
+
None, #推理不需要NUM_PATCHES
|
| 1270 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1271 |
+
action_head,
|
| 1272 |
+
)
|
| 1273 |
+
|
| 1274 |
+
# Unnormalize predicted actions
|
| 1275 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1276 |
+
|
| 1277 |
+
return actions, actions_hidden_states
|
| 1278 |
+
|
| 1279 |
+
|
| 1280 |
+
def mul_predict_action(
|
| 1281 |
+
self,
|
| 1282 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1283 |
+
unnorm_key: Optional[str] = None,
|
| 1284 |
+
proprio=None,
|
| 1285 |
+
proprio_projector=None,
|
| 1286 |
+
action_head:L1RegressionActionHead=None,
|
| 1287 |
+
noisy_action_projector=None,
|
| 1288 |
+
use_film: bool = False,
|
| 1289 |
+
**kwargs: str,
|
| 1290 |
+
) -> np.ndarray:
|
| 1291 |
+
|
| 1292 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 1293 |
+
input_ids = torch.cat(
|
| 1294 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 1295 |
+
)
|
| 1296 |
+
|
| 1297 |
+
pixel_values = kwargs["pixel_values"]
|
| 1298 |
+
attention_mask = kwargs["attention_mask"]
|
| 1299 |
+
|
| 1300 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1301 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1302 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1303 |
+
|
| 1304 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1305 |
+
|
| 1306 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1307 |
+
if use_proprio:
|
| 1308 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1309 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1310 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1311 |
+
)
|
| 1312 |
+
|
| 1313 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1314 |
+
input_embeddings,
|
| 1315 |
+
None,
|
| 1316 |
+
projected_patch_embeddings,
|
| 1317 |
+
attention_mask,
|
| 1318 |
+
None, #推理不需要labels
|
| 1319 |
+
None, #推理不需要NUM_PATCHES
|
| 1320 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1321 |
+
action_head,
|
| 1322 |
+
)
|
| 1323 |
+
|
| 1324 |
+
# Unnormalize predicted actions
|
| 1325 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1326 |
+
|
| 1327 |
+
return actions, actions_hidden_states
|
| 1328 |
+
|
modeling_prismatic.py.back.20250403_130245
ADDED
|
@@ -0,0 +1,1328 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
# 0 , is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 957 |
+
# actions_hidden_states = language_model_output.hidden_states[1][-1]
|
| 958 |
+
import ipdb; ipdb.set_trace()
|
| 959 |
+
actions_hidden_states = language_model_output.hidden_states[0][-1]
|
| 960 |
+
actions_hidden_states = actions_hidden_states[:, -1]
|
| 961 |
+
|
| 962 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 963 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 964 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 965 |
+
|
| 966 |
+
return normalized_actions, actions_hidden_states
|
| 967 |
+
|
| 968 |
+
def predict_action(
|
| 969 |
+
self,
|
| 970 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 971 |
+
unnorm_key: Optional[str] = None,
|
| 972 |
+
proprio=None,
|
| 973 |
+
proprio_projector=None,
|
| 974 |
+
action_head=None,
|
| 975 |
+
noisy_action_projector=None,
|
| 976 |
+
use_film: bool = False,
|
| 977 |
+
**kwargs: str,
|
| 978 |
+
) -> np.ndarray:
|
| 979 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 980 |
+
|
| 981 |
+
Args:
|
| 982 |
+
input_ids: Input token ids
|
| 983 |
+
unnorm_key: Key for unnormalization statistics
|
| 984 |
+
proprio: Proprioceptive features
|
| 985 |
+
proprio_projector: Projector for proprioceptive features
|
| 986 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 987 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 988 |
+
use_film: Whether to use FiLM conditioning
|
| 989 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 990 |
+
|
| 991 |
+
Returns:
|
| 992 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 993 |
+
"""
|
| 994 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 995 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 996 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 997 |
+
input_ids = torch.cat(
|
| 998 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 999 |
+
)
|
| 1000 |
+
|
| 1001 |
+
pixel_values = kwargs["pixel_values"]
|
| 1002 |
+
attention_mask = kwargs["attention_mask"]
|
| 1003 |
+
|
| 1004 |
+
# Create fake labels tensor (needed for action mask)
|
| 1005 |
+
labels = input_ids.clone()
|
| 1006 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1007 |
+
|
| 1008 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1009 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1010 |
+
|
| 1011 |
+
# Prepare inputs by adding necessary tokens
|
| 1012 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1013 |
+
|
| 1014 |
+
# Update labels tensor for action mask computation later
|
| 1015 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1016 |
+
|
| 1017 |
+
# Get input embeddings and action masks
|
| 1018 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1019 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1020 |
+
|
| 1021 |
+
# Extract language embeddings
|
| 1022 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1023 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1024 |
+
)
|
| 1025 |
+
|
| 1026 |
+
# Process vision features
|
| 1027 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1028 |
+
|
| 1029 |
+
# Add proprioceptive features if provided
|
| 1030 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1031 |
+
if use_proprio:
|
| 1032 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1033 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1034 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1035 |
+
)
|
| 1036 |
+
|
| 1037 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1038 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1039 |
+
|
| 1040 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1041 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1042 |
+
if use_proprio:
|
| 1043 |
+
NUM_PATCHES += 1
|
| 1044 |
+
|
| 1045 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1046 |
+
input_embeddings,
|
| 1047 |
+
all_actions_mask,
|
| 1048 |
+
projected_patch_embeddings,
|
| 1049 |
+
attention_mask,
|
| 1050 |
+
labels,
|
| 1051 |
+
NUM_PATCHES,
|
| 1052 |
+
NUM_PROMPT_TOKENS,
|
| 1053 |
+
action_head,
|
| 1054 |
+
)
|
| 1055 |
+
|
| 1056 |
+
# Unnormalize predicted actions
|
| 1057 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1058 |
+
|
| 1059 |
+
return actions, actions_hidden_states
|
| 1060 |
+
|
| 1061 |
+
@staticmethod
|
| 1062 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1063 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1064 |
+
if unnorm_key is None:
|
| 1065 |
+
assert len(norm_stats) == 1, (
|
| 1066 |
+
f"Your model was trained on more than one dataset, "
|
| 1067 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1068 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1069 |
+
)
|
| 1070 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1071 |
+
# norm states没有加载libero, 为什么?
|
| 1072 |
+
assert unnorm_key in norm_stats, (
|
| 1073 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1074 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1075 |
+
)
|
| 1076 |
+
return unnorm_key
|
| 1077 |
+
|
| 1078 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1079 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1080 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1081 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1082 |
+
|
| 1083 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1084 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1085 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1086 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1087 |
+
|
| 1088 |
+
|
| 1089 |
+
def lisa_forward(
|
| 1090 |
+
self,
|
| 1091 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1092 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1093 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1094 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1095 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1096 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1097 |
+
use_cache: Optional[bool] = None,
|
| 1098 |
+
output_attentions: Optional[bool] = None,
|
| 1099 |
+
output_hidden_states: Optional[bool] = None,
|
| 1100 |
+
output_projector_features: Optional[bool] = None,
|
| 1101 |
+
return_dict: Optional[bool] = None,
|
| 1102 |
+
proprio=None,
|
| 1103 |
+
proprio_projector=None,
|
| 1104 |
+
**kwargs
|
| 1105 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1106 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1107 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1108 |
+
output_hidden_states = (
|
| 1109 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1110 |
+
)
|
| 1111 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1112 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1113 |
+
|
| 1114 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1115 |
+
use_cache = use_cache and not self.training
|
| 1116 |
+
|
| 1117 |
+
# Instantiate Placeholder for Projector Features
|
| 1118 |
+
projected_patch_embeddings = None
|
| 1119 |
+
|
| 1120 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1121 |
+
if input_ids.shape[1] == 1:
|
| 1122 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1123 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1124 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1125 |
+
|
| 1126 |
+
language_model_output = self.language_model(
|
| 1127 |
+
input_ids=input_ids,
|
| 1128 |
+
attention_mask=None,
|
| 1129 |
+
position_ids=None,
|
| 1130 |
+
past_key_values=past_key_values,
|
| 1131 |
+
inputs_embeds=None,
|
| 1132 |
+
labels=None,
|
| 1133 |
+
use_cache=use_cache,
|
| 1134 |
+
output_attentions=output_attentions,
|
| 1135 |
+
output_hidden_states=output_hidden_states,
|
| 1136 |
+
return_dict=return_dict,
|
| 1137 |
+
)
|
| 1138 |
+
|
| 1139 |
+
# === Handle Unimodal Forward ===
|
| 1140 |
+
elif pixel_values is None:
|
| 1141 |
+
raise NotImplementedError
|
| 1142 |
+
|
| 1143 |
+
# === Handle Multimodal Forward ===
|
| 1144 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1145 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1146 |
+
|
| 1147 |
+
# Get input embeddings (from language model embeddings)
|
| 1148 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1149 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1150 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1151 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1152 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1153 |
+
# Get visual features
|
| 1154 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1155 |
+
|
| 1156 |
+
# Add proprioceptive state if provided
|
| 1157 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1158 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1159 |
+
)
|
| 1160 |
+
|
| 1161 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1162 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1163 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1164 |
+
|
| 1165 |
+
# Build multimodal embeddings & attention mask
|
| 1166 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1167 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1168 |
+
)
|
| 1169 |
+
|
| 1170 |
+
# Build labels for multimodal sequence if needed
|
| 1171 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1172 |
+
|
| 1173 |
+
# Dispatch to language model
|
| 1174 |
+
language_model_output = self.language_model(
|
| 1175 |
+
input_ids=None,
|
| 1176 |
+
attention_mask=multimodal_attention_mask,
|
| 1177 |
+
position_ids=None,
|
| 1178 |
+
past_key_values=None,
|
| 1179 |
+
inputs_embeds=multimodal_embeddings,
|
| 1180 |
+
labels=multimodal_labels,
|
| 1181 |
+
use_cache=use_cache,
|
| 1182 |
+
output_attentions=output_attentions,
|
| 1183 |
+
output_hidden_states=output_hidden_states,
|
| 1184 |
+
return_dict=return_dict,
|
| 1185 |
+
)
|
| 1186 |
+
|
| 1187 |
+
|
| 1188 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1189 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1190 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1191 |
+
|
| 1192 |
+
else:
|
| 1193 |
+
raise ValueError(
|
| 1194 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1195 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1196 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1197 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1198 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1199 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1200 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1201 |
+
f"=> `use_cache` = {use_cache}"
|
| 1202 |
+
)
|
| 1203 |
+
|
| 1204 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1205 |
+
if not return_dict:
|
| 1206 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1207 |
+
return *language_model_output, projected_patch_embeddings
|
| 1208 |
+
|
| 1209 |
+
return language_model_output
|
| 1210 |
+
|
| 1211 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1212 |
+
loss=language_model_output.loss,
|
| 1213 |
+
logits=language_model_output.logits,
|
| 1214 |
+
past_key_values=language_model_output.past_key_values,
|
| 1215 |
+
hidden_states=language_model_output.hidden_states,
|
| 1216 |
+
attentions=language_model_output.attentions,
|
| 1217 |
+
projector_features=projected_patch_embeddings,
|
| 1218 |
+
)
|
| 1219 |
+
|
| 1220 |
+
def lisa_predict_action(
|
| 1221 |
+
self,
|
| 1222 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1223 |
+
unnorm_key: Optional[str] = None,
|
| 1224 |
+
proprio=None,
|
| 1225 |
+
proprio_projector=None,
|
| 1226 |
+
action_head:L1RegressionActionHead=None,
|
| 1227 |
+
noisy_action_projector=None,
|
| 1228 |
+
use_film: bool = False,
|
| 1229 |
+
**kwargs: str,
|
| 1230 |
+
) -> np.ndarray:
|
| 1231 |
+
|
| 1232 |
+
pixel_values = kwargs["pixel_values"]
|
| 1233 |
+
attention_mask = kwargs["attention_mask"]
|
| 1234 |
+
|
| 1235 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1236 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1237 |
+
|
| 1238 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1239 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1240 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1241 |
+
if cfg.preset:
|
| 1242 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1243 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1244 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1245 |
+
self.preset = True
|
| 1246 |
+
else:
|
| 1247 |
+
self.preset = False
|
| 1248 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1249 |
+
|
| 1250 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1251 |
+
|
| 1252 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1253 |
+
if use_proprio:
|
| 1254 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1255 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1256 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1257 |
+
)
|
| 1258 |
+
|
| 1259 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1260 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1261 |
+
# if use_proprio:
|
| 1262 |
+
# NUM_PATCHES += 1
|
| 1263 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1264 |
+
input_embeddings,
|
| 1265 |
+
None,
|
| 1266 |
+
projected_patch_embeddings,
|
| 1267 |
+
attention_mask,
|
| 1268 |
+
None, #推理不需要labels
|
| 1269 |
+
None, #推理不需要NUM_PATCHES
|
| 1270 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1271 |
+
action_head,
|
| 1272 |
+
)
|
| 1273 |
+
|
| 1274 |
+
# Unnormalize predicted actions
|
| 1275 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1276 |
+
|
| 1277 |
+
return actions, actions_hidden_states
|
| 1278 |
+
|
| 1279 |
+
|
| 1280 |
+
def mul_predict_action(
|
| 1281 |
+
self,
|
| 1282 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1283 |
+
unnorm_key: Optional[str] = None,
|
| 1284 |
+
proprio=None,
|
| 1285 |
+
proprio_projector=None,
|
| 1286 |
+
action_head:L1RegressionActionHead=None,
|
| 1287 |
+
noisy_action_projector=None,
|
| 1288 |
+
use_film: bool = False,
|
| 1289 |
+
**kwargs: str,
|
| 1290 |
+
) -> np.ndarray:
|
| 1291 |
+
|
| 1292 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 1293 |
+
input_ids = torch.cat(
|
| 1294 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 1295 |
+
)
|
| 1296 |
+
|
| 1297 |
+
pixel_values = kwargs["pixel_values"]
|
| 1298 |
+
attention_mask = kwargs["attention_mask"]
|
| 1299 |
+
|
| 1300 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1301 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1302 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1303 |
+
|
| 1304 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1305 |
+
|
| 1306 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1307 |
+
if use_proprio:
|
| 1308 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1309 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1310 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1311 |
+
)
|
| 1312 |
+
|
| 1313 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1314 |
+
input_embeddings,
|
| 1315 |
+
None,
|
| 1316 |
+
projected_patch_embeddings,
|
| 1317 |
+
attention_mask,
|
| 1318 |
+
None, #推理不需要labels
|
| 1319 |
+
None, #推理不需要NUM_PATCHES
|
| 1320 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1321 |
+
action_head,
|
| 1322 |
+
)
|
| 1323 |
+
|
| 1324 |
+
# Unnormalize predicted actions
|
| 1325 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1326 |
+
|
| 1327 |
+
return actions, actions_hidden_states
|
| 1328 |
+
|
modeling_prismatic.py.back.20250403_132953
ADDED
|
@@ -0,0 +1,1327 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
from dataclasses import dataclass
|
| 11 |
+
from functools import partial
|
| 12 |
+
from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import timm
|
| 16 |
+
import tokenizers
|
| 17 |
+
import torch
|
| 18 |
+
import torch.nn as nn
|
| 19 |
+
import transformers
|
| 20 |
+
from timm.models.vision_transformer import LayerScale
|
| 21 |
+
from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
|
| 22 |
+
from transformers.modeling_outputs import ModelOutput
|
| 23 |
+
from prismatic.models.action_heads import L1RegressionActionHead
|
| 24 |
+
import time
|
| 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 |
+
ACTION_TOKEN_IDX,
|
| 37 |
+
NormalizationType,
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
|
| 41 |
+
|
| 42 |
+
# Set up logger
|
| 43 |
+
logger = logging.getLogger(__name__)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
# === Utility Functions for Monkey-Patching ===
|
| 47 |
+
def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
|
| 48 |
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
| 49 |
+
result = fn(*args, **kwargs)
|
| 50 |
+
return result[0] if isinstance(result, tuple) else result
|
| 51 |
+
|
| 52 |
+
return wrapper
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
|
| 56 |
+
# =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
|
| 57 |
+
# =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
|
| 58 |
+
def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 59 |
+
return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def ls_apply_patch(ls_module: LayerScale):
|
| 63 |
+
ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
|
| 64 |
+
ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
|
| 65 |
+
del ls_module.gamma
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
|
| 69 |
+
class PrismaticVisionBackbone(nn.Module):
|
| 70 |
+
"""
|
| 71 |
+
Vision backbone for Prismatic models that handles image feature extraction.
|
| 72 |
+
|
| 73 |
+
Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
|
| 74 |
+
For fused backbones, features from both models are concatenated along the feature dimension.
|
| 75 |
+
"""
|
| 76 |
+
|
| 77 |
+
def __init__(
|
| 78 |
+
self,
|
| 79 |
+
use_fused_vision_backbone: bool,
|
| 80 |
+
image_sizes: List[int],
|
| 81 |
+
timm_model_ids: List[str],
|
| 82 |
+
timm_override_act_layers: List[Optional[str]],
|
| 83 |
+
) -> None:
|
| 84 |
+
"""
|
| 85 |
+
Initialize the vision backbone.
|
| 86 |
+
|
| 87 |
+
Args:
|
| 88 |
+
use_fused_vision_backbone: Whether to use two backbones and fuse their features
|
| 89 |
+
image_sizes: List of image sizes for each backbone
|
| 90 |
+
timm_model_ids: List of TIMM model IDs to use for each backbone
|
| 91 |
+
timm_override_act_layers: List of activation layer overrides for each backbone
|
| 92 |
+
"""
|
| 93 |
+
super().__init__()
|
| 94 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 95 |
+
self.num_images_in_input = 1 # Default value, can be overridden later
|
| 96 |
+
|
| 97 |
+
# Validate number of (fused) vision backbones
|
| 98 |
+
if len(timm_model_ids) > 2:
|
| 99 |
+
raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
|
| 100 |
+
|
| 101 |
+
# Create primary featurizer
|
| 102 |
+
self.featurizer = self._create_featurizer(
|
| 103 |
+
model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
|
| 104 |
+
)
|
| 105 |
+
self.embed_dim = self.featurizer.embed_dim
|
| 106 |
+
|
| 107 |
+
# Create secondary featurizer if using fused backbone
|
| 108 |
+
if self.use_fused_vision_backbone:
|
| 109 |
+
self.fused_featurizer = self._create_featurizer(
|
| 110 |
+
model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
|
| 111 |
+
)
|
| 112 |
+
self.embed_dim += self.fused_featurizer.embed_dim
|
| 113 |
+
|
| 114 |
+
# Patch LayerScale modules for HF compatibility
|
| 115 |
+
self._patch_layer_scales()
|
| 116 |
+
|
| 117 |
+
def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
|
| 118 |
+
"""
|
| 119 |
+
Create a TIMM-based featurizer model with appropriate configurations.
|
| 120 |
+
|
| 121 |
+
Args:
|
| 122 |
+
model_id: The TIMM model ID to load
|
| 123 |
+
img_size: Input image size for the model
|
| 124 |
+
act_layer: Override for the activation layer type
|
| 125 |
+
|
| 126 |
+
Returns:
|
| 127 |
+
A configured featurizer model
|
| 128 |
+
"""
|
| 129 |
+
featurizer = timm.create_model(
|
| 130 |
+
model_id,
|
| 131 |
+
pretrained=False,
|
| 132 |
+
num_classes=0,
|
| 133 |
+
img_size=img_size,
|
| 134 |
+
act_layer=act_layer,
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
# Monkey-patch the forward function to extract the second-to-last layer features
|
| 138 |
+
num_blocks = len(featurizer.blocks)
|
| 139 |
+
featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
|
| 140 |
+
|
| 141 |
+
return featurizer
|
| 142 |
+
|
| 143 |
+
def _patch_layer_scales(self) -> None:
|
| 144 |
+
"""
|
| 145 |
+
Patch all LayerScale modules to be compatible with HF's parameter naming.
|
| 146 |
+
|
| 147 |
+
HF Transformers overwrites parameters with names containing 'gamma',
|
| 148 |
+
so we need to rename and modify the forward method.
|
| 149 |
+
"""
|
| 150 |
+
# Patch primary featurizer
|
| 151 |
+
for module in self.featurizer.modules():
|
| 152 |
+
if isinstance(module, LayerScale):
|
| 153 |
+
ls_apply_patch(module)
|
| 154 |
+
|
| 155 |
+
# Patch secondary featurizer if it exists
|
| 156 |
+
if self.use_fused_vision_backbone:
|
| 157 |
+
for module in self.fused_featurizer.modules():
|
| 158 |
+
if isinstance(module, LayerScale):
|
| 159 |
+
ls_apply_patch(module)
|
| 160 |
+
|
| 161 |
+
def get_num_patches(self) -> int:
|
| 162 |
+
"""
|
| 163 |
+
Returns the number of vision patches output by the vision backbone.
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
Number of patches per image
|
| 167 |
+
"""
|
| 168 |
+
return self.featurizer.patch_embed.num_patches
|
| 169 |
+
|
| 170 |
+
def get_num_images_in_input(self) -> int:
|
| 171 |
+
"""
|
| 172 |
+
Returns the number of input images for the vision backbone.
|
| 173 |
+
|
| 174 |
+
Returns:
|
| 175 |
+
Number of images expected in the input
|
| 176 |
+
"""
|
| 177 |
+
return self.num_images_in_input
|
| 178 |
+
|
| 179 |
+
def set_num_images_in_input(self, num_images_in_input: int) -> None:
|
| 180 |
+
"""
|
| 181 |
+
Sets the number of input images for the vision backbone.
|
| 182 |
+
|
| 183 |
+
Args:
|
| 184 |
+
num_images_in_input: Number of images to expect in the input
|
| 185 |
+
"""
|
| 186 |
+
self.num_images_in_input = num_images_in_input
|
| 187 |
+
|
| 188 |
+
def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
|
| 189 |
+
"""
|
| 190 |
+
Implements the forward pass for the vision backbone.
|
| 191 |
+
|
| 192 |
+
If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
|
| 193 |
+
(otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
|
| 194 |
+
|
| 195 |
+
Args:
|
| 196 |
+
pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
|
| 197 |
+
"""
|
| 198 |
+
if self.num_images_in_input == 1:
|
| 199 |
+
if not self.use_fused_vision_backbone:
|
| 200 |
+
return self.featurizer(pixel_values)
|
| 201 |
+
|
| 202 |
+
# Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
|
| 203 |
+
img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
|
| 204 |
+
patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
|
| 205 |
+
|
| 206 |
+
return torch.cat([patches, patches_fused], dim=2)
|
| 207 |
+
|
| 208 |
+
else:
|
| 209 |
+
assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
|
| 210 |
+
|
| 211 |
+
# Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
|
| 212 |
+
images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
|
| 213 |
+
|
| 214 |
+
# Process each image and collect patches
|
| 215 |
+
all_patches = []
|
| 216 |
+
for img in images:
|
| 217 |
+
# Split each image further into two stacks of channels (each with 3 channels)
|
| 218 |
+
img_regular, img_fused = torch.split(img, [3, 3], dim=1)
|
| 219 |
+
|
| 220 |
+
# Get patches from both SigLIP and DINOv2 vision transformers
|
| 221 |
+
patches = self.featurizer(img_regular)
|
| 222 |
+
patches_fused = self.fused_featurizer(img_fused)
|
| 223 |
+
|
| 224 |
+
# Concatenate SigLIP and DINOv2 patches along the hidden dimension
|
| 225 |
+
combined_patches = torch.cat([patches, patches_fused], dim=2)
|
| 226 |
+
all_patches.append(combined_patches)
|
| 227 |
+
|
| 228 |
+
# Concatenate all patches along the patch dimension
|
| 229 |
+
return torch.cat(all_patches, dim=1)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
# === Prismatic Projector (nn.Module) Definitions ===
|
| 233 |
+
class PrismaticProjector(nn.Module):
|
| 234 |
+
def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
|
| 235 |
+
super().__init__()
|
| 236 |
+
self.use_fused_vision_backbone = use_fused_vision_backbone
|
| 237 |
+
self.vision_dim, self.llm_dim = vision_dim, llm_dim
|
| 238 |
+
|
| 239 |
+
# Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
|
| 240 |
+
if not self.use_fused_vision_backbone:
|
| 241 |
+
self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
|
| 242 |
+
self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 243 |
+
self.act_fn1 = nn.GELU()
|
| 244 |
+
else:
|
| 245 |
+
initial_projection_dim = 4 * vision_dim
|
| 246 |
+
self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
|
| 247 |
+
self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
|
| 248 |
+
self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
|
| 249 |
+
self.act_fn1 = nn.GELU()
|
| 250 |
+
self.act_fn2 = nn.GELU()
|
| 251 |
+
|
| 252 |
+
def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
|
| 253 |
+
if not self.use_fused_vision_backbone:
|
| 254 |
+
projected_features = self.fc1(img_patches)
|
| 255 |
+
projected_features = self.act_fn1(projected_features)
|
| 256 |
+
projected_features = self.fc2(projected_features)
|
| 257 |
+
else:
|
| 258 |
+
projected_features = self.fc1(img_patches)
|
| 259 |
+
projected_features = self.act_fn1(projected_features)
|
| 260 |
+
projected_features = self.fc2(projected_features)
|
| 261 |
+
projected_features = self.act_fn2(projected_features)
|
| 262 |
+
projected_features = self.fc3(projected_features)
|
| 263 |
+
|
| 264 |
+
return projected_features
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
# === Main HF Class Definitions ===
|
| 268 |
+
@dataclass
|
| 269 |
+
class PrismaticCausalLMOutputWithPast(ModelOutput):
|
| 270 |
+
"""Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
|
| 271 |
+
|
| 272 |
+
loss: Optional[torch.FloatTensor] = None
|
| 273 |
+
logits: torch.FloatTensor = None
|
| 274 |
+
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
|
| 275 |
+
hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
|
| 276 |
+
attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| 277 |
+
|
| 278 |
+
# Additions for VLMs
|
| 279 |
+
projector_features: Optional[torch.FloatTensor] = None
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
class PrismaticPreTrainedModel(PreTrainedModel):
|
| 283 |
+
config_class: PretrainedConfig = PrismaticConfig
|
| 284 |
+
base_model_prefix: str = "model"
|
| 285 |
+
supports_gradient_checkpointing: bool = True
|
| 286 |
+
|
| 287 |
+
_no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
|
| 288 |
+
_skip_keys_device_placement: str = "past_key_values"
|
| 289 |
+
_supports_flash_attn_2: bool = True
|
| 290 |
+
|
| 291 |
+
def _init_weights(self, module: nn.Module) -> None:
|
| 292 |
+
# Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
|
| 293 |
+
# => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
|
| 294 |
+
# https://github.com/TRI-ML/prismatic-vlms
|
| 295 |
+
std = (
|
| 296 |
+
self.config.initializer_range
|
| 297 |
+
if hasattr(self.config, "initializer_range")
|
| 298 |
+
else self.config.text_config.initializer_range
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
if hasattr(module, "class_embedding"):
|
| 302 |
+
module.class_embedding.data.normal_(mean=0.0, std=std)
|
| 303 |
+
|
| 304 |
+
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
| 305 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 306 |
+
if module.bias is not None:
|
| 307 |
+
module.bias.data.zero_()
|
| 308 |
+
elif isinstance(module, nn.Embedding):
|
| 309 |
+
module.weight.data.normal_(mean=0.0, std=std)
|
| 310 |
+
if module.padding_idx is not None:
|
| 311 |
+
module.weight.data[module.padding_idx].zero_()
|
| 312 |
+
|
| 313 |
+
@property
|
| 314 |
+
def _supports_sdpa(self) -> bool:
|
| 315 |
+
"""Check LLM supports SDPA Attention"""
|
| 316 |
+
return self.language_model._supports_sdpa
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
|
| 320 |
+
def __init__(self, config: PrismaticConfig) -> None:
|
| 321 |
+
super().__init__(config)
|
| 322 |
+
|
| 323 |
+
# [Validation] Lightweight Validate on `config` Fields + Dependency Versions
|
| 324 |
+
if config.use_fused_vision_backbone is None:
|
| 325 |
+
raise ValueError("Missing config field `use_fused_vision_backbone`")
|
| 326 |
+
|
| 327 |
+
if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
|
| 328 |
+
raise NotImplementedError(
|
| 329 |
+
"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
|
| 330 |
+
"if you urgently need support for latest TIMM versions."
|
| 331 |
+
)
|
| 332 |
+
|
| 333 |
+
if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
|
| 334 |
+
logger.warning(
|
| 335 |
+
f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
|
| 336 |
+
f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
|
| 337 |
+
f"there might be inference-time regressions due to dependency changes. If in doubt, please"
|
| 338 |
+
f"use the above versions."
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
# Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
|
| 342 |
+
self.vision_backbone = PrismaticVisionBackbone(
|
| 343 |
+
config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
|
| 344 |
+
)
|
| 345 |
+
|
| 346 |
+
# Create Multimodal Projector
|
| 347 |
+
self.projector = PrismaticProjector(
|
| 348 |
+
config.use_fused_vision_backbone,
|
| 349 |
+
vision_dim=self.vision_backbone.embed_dim,
|
| 350 |
+
llm_dim=config.text_config.hidden_size,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
+
# Instantiate LLM Backbone
|
| 354 |
+
self.language_model = AutoModelForCausalLM.from_config(
|
| 355 |
+
config.text_config, attn_implementation=config._attn_implementation
|
| 356 |
+
)
|
| 357 |
+
self.vocab_size = config.text_config.vocab_size
|
| 358 |
+
self.pad_token_id = config.pad_token_id
|
| 359 |
+
self.llm_dim = config.text_config.hidden_size
|
| 360 |
+
|
| 361 |
+
# HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
|
| 362 |
+
self.post_init()
|
| 363 |
+
|
| 364 |
+
# === `PreTrainedModel` Boilerplate ===
|
| 365 |
+
def get_input_embeddings(self) -> nn.Module:
|
| 366 |
+
return self.language_model.get_input_embeddings()
|
| 367 |
+
|
| 368 |
+
def set_input_embeddings(self, value: nn.Module) -> None:
|
| 369 |
+
self.language_model.set_input_embeddings(value)
|
| 370 |
+
|
| 371 |
+
def get_output_embeddings(self) -> nn.Module:
|
| 372 |
+
return self.language_model.get_output_embeddings()
|
| 373 |
+
|
| 374 |
+
def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
|
| 375 |
+
self.language_model.set_output_embeddings(new_embeddings)
|
| 376 |
+
|
| 377 |
+
def get_decoder(self) -> nn.Module:
|
| 378 |
+
return self.language_model.get_decoder()
|
| 379 |
+
|
| 380 |
+
def set_decoder(self, decoder: nn.Module) -> None:
|
| 381 |
+
self.language_model.set_decoder(decoder)
|
| 382 |
+
|
| 383 |
+
def tie_weights(self) -> None:
|
| 384 |
+
self.language_model.tie_weights() # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
|
| 385 |
+
|
| 386 |
+
def resize_token_embeddings(
|
| 387 |
+
self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
|
| 388 |
+
) -> nn.Embedding:
|
| 389 |
+
updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
|
| 390 |
+
|
| 391 |
+
# Update config/instance variables
|
| 392 |
+
self.config.text_config.vocab_size = updated_embeddings.num_embeddings
|
| 393 |
+
self.vocab_size = updated_embeddings.num_embeddings
|
| 394 |
+
|
| 395 |
+
return updated_embeddings
|
| 396 |
+
|
| 397 |
+
def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
|
| 398 |
+
"""
|
| 399 |
+
Replace embeddings in input_embeddings at positions where all_actions_mask is True
|
| 400 |
+
with embeddings from noisy_action_features, using vectorized operations.
|
| 401 |
+
|
| 402 |
+
Args:
|
| 403 |
+
input_embeddings: Tensor of shape (B, S, D)
|
| 404 |
+
all_actions_mask: Boolean tensor of shape (B, S)
|
| 405 |
+
noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
|
| 406 |
+
|
| 407 |
+
Returns:
|
| 408 |
+
Modified input_embeddings tensor
|
| 409 |
+
"""
|
| 410 |
+
# Clone input to avoid modifying the original tensor
|
| 411 |
+
new_input_embeddings = input_embeddings.clone()
|
| 412 |
+
|
| 413 |
+
# Create a tensor with the same shape of input_embeddings to hold the noisy action features
|
| 414 |
+
repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
|
| 415 |
+
|
| 416 |
+
# Create batch indices for splicing
|
| 417 |
+
batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
|
| 418 |
+
batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
|
| 419 |
+
|
| 420 |
+
# Get indices where mask is True for each sample
|
| 421 |
+
masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
|
| 422 |
+
|
| 423 |
+
# Move the noisy action features into their correct positions
|
| 424 |
+
repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
|
| 425 |
+
|
| 426 |
+
# Combine original input embeddings and noisy action embeddings using the mask
|
| 427 |
+
new_input_embeddings = torch.where(
|
| 428 |
+
all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
|
| 429 |
+
)
|
| 430 |
+
|
| 431 |
+
return new_input_embeddings
|
| 432 |
+
|
| 433 |
+
def _process_action_masks(self, labels):
|
| 434 |
+
"""Helper to get action masks from labels"""
|
| 435 |
+
current_action_mask = get_current_action_mask(labels) # (B, seq_len)
|
| 436 |
+
next_actions_mask = get_next_actions_mask(labels) # (B, seq_len)
|
| 437 |
+
all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len)
|
| 438 |
+
return all_actions_mask
|
| 439 |
+
|
| 440 |
+
def _process_vision_features(self, pixel_values):
|
| 441 |
+
"""Process vision features with optional FiLM conditioning"""
|
| 442 |
+
patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D)
|
| 443 |
+
|
| 444 |
+
# Project patch embeddings into language embedding space
|
| 445 |
+
return self.projector(patch_features)
|
| 446 |
+
|
| 447 |
+
def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
|
| 448 |
+
"""Process proprioceptive features and append to vision features"""
|
| 449 |
+
if proprio_projector is not None and proprio is not None:
|
| 450 |
+
# projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
|
| 451 |
+
# proprio: (bsz, proprio_dim) or (propro_dim,)
|
| 452 |
+
proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim)
|
| 453 |
+
proprio_features = proprio_projector(proprio) # (bsz, llm_dim)
|
| 454 |
+
proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim)
|
| 455 |
+
# For simplicity, just append proprio token to the end of projected vision patch tokens
|
| 456 |
+
return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
|
| 457 |
+
return projected_patch_embeddings
|
| 458 |
+
|
| 459 |
+
def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
|
| 460 |
+
"""Build multimodal embeddings and attention mask"""
|
| 461 |
+
# juyi: Update attention mask 是不是要改成下三角? 不用, 因为generate会自动屏蔽
|
| 462 |
+
projected_patch_attention_mask = None
|
| 463 |
+
if attention_mask is not None:
|
| 464 |
+
projected_patch_attention_mask = torch.full(
|
| 465 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 466 |
+
fill_value=True,
|
| 467 |
+
dtype=attention_mask.dtype,
|
| 468 |
+
device=attention_mask.device,
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
# Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
|
| 472 |
+
multimodal_embeddings = torch.cat(
|
| 473 |
+
[input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
|
| 474 |
+
)
|
| 475 |
+
|
| 476 |
+
multimodal_attention_mask = None
|
| 477 |
+
if attention_mask is not None:
|
| 478 |
+
multimodal_attention_mask = torch.cat(
|
| 479 |
+
[attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
|
| 480 |
+
)
|
| 481 |
+
|
| 482 |
+
return multimodal_embeddings, multimodal_attention_mask
|
| 483 |
+
|
| 484 |
+
def _build_multimodal_labels(self, labels, projected_patch_embeddings):
|
| 485 |
+
"""Build multimodal labels with IGNORE_INDEX for patch embeddings"""
|
| 486 |
+
if labels is not None:
|
| 487 |
+
projected_patch_labels = torch.full(
|
| 488 |
+
(projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
|
| 489 |
+
fill_value=IGNORE_INDEX, # 这些位置不需要计算损失。
|
| 490 |
+
dtype=labels.dtype,
|
| 491 |
+
device=labels.device,
|
| 492 |
+
)
|
| 493 |
+
return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) # 第一个token是<BOS>
|
| 494 |
+
return None
|
| 495 |
+
|
| 496 |
+
# === Core Prismatic VLM `forward()` Logic ===
|
| 497 |
+
def forward(
|
| 498 |
+
self,
|
| 499 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 500 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 501 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 502 |
+
labels: Optional[torch.LongTensor] = None,
|
| 503 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 504 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 505 |
+
use_cache: Optional[bool] = None,
|
| 506 |
+
output_attentions: Optional[bool] = None,
|
| 507 |
+
output_hidden_states: Optional[bool] = None,
|
| 508 |
+
output_projector_features: Optional[bool] = None,
|
| 509 |
+
return_dict: Optional[bool] = None,
|
| 510 |
+
proprio=None,
|
| 511 |
+
proprio_projector=None,
|
| 512 |
+
noisy_actions=None,
|
| 513 |
+
noisy_action_projector=None,
|
| 514 |
+
diffusion_timestep_embeddings=None,
|
| 515 |
+
use_film: bool = False,
|
| 516 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 517 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 518 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 519 |
+
output_hidden_states = (
|
| 520 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 521 |
+
)
|
| 522 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 523 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 524 |
+
|
| 525 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 526 |
+
use_cache = use_cache and not self.training
|
| 527 |
+
|
| 528 |
+
# Instantiate Placeholder for Projector Features
|
| 529 |
+
projected_patch_embeddings = None
|
| 530 |
+
|
| 531 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 532 |
+
if input_ids.shape[1] == 1:
|
| 533 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 534 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 535 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 536 |
+
|
| 537 |
+
language_model_output = self.language_model(
|
| 538 |
+
input_ids=input_ids,
|
| 539 |
+
attention_mask=None,
|
| 540 |
+
position_ids=None,
|
| 541 |
+
past_key_values=past_key_values,
|
| 542 |
+
inputs_embeds=None,
|
| 543 |
+
labels=None,
|
| 544 |
+
use_cache=use_cache,
|
| 545 |
+
output_attentions=output_attentions,
|
| 546 |
+
output_hidden_states=output_hidden_states,
|
| 547 |
+
return_dict=return_dict,
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# === Handle Unimodal Forward ===
|
| 551 |
+
elif pixel_values is None:
|
| 552 |
+
assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
|
| 553 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
|
| 554 |
+
|
| 555 |
+
language_model_output = self.language_model(
|
| 556 |
+
input_ids=input_ids,
|
| 557 |
+
attention_mask=attention_mask,
|
| 558 |
+
position_ids=None,
|
| 559 |
+
past_key_values=None,
|
| 560 |
+
inputs_embeds=None,
|
| 561 |
+
labels=labels,
|
| 562 |
+
use_cache=use_cache,
|
| 563 |
+
output_attentions=output_attentions,
|
| 564 |
+
output_hidden_states=output_hidden_states,
|
| 565 |
+
return_dict=return_dict,
|
| 566 |
+
)
|
| 567 |
+
|
| 568 |
+
# === Handle Multimodal Forward ===
|
| 569 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 570 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 571 |
+
|
| 572 |
+
# Get input embeddings (from language model embeddings)
|
| 573 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 574 |
+
|
| 575 |
+
# Extract action masks
|
| 576 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 577 |
+
|
| 578 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 579 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 580 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 581 |
+
) # (B, lang_seq_len, llm_dim)
|
| 582 |
+
|
| 583 |
+
# Get visual features
|
| 584 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 585 |
+
|
| 586 |
+
# Add proprioceptive state if provided
|
| 587 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 588 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 589 |
+
)
|
| 590 |
+
|
| 591 |
+
# [Diffusion] Add diffusion timestep embedding if provided
|
| 592 |
+
if diffusion_timestep_embeddings is not None:
|
| 593 |
+
# For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
|
| 594 |
+
projected_patch_embeddings = torch.cat(
|
| 595 |
+
(projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 596 |
+
)
|
| 597 |
+
|
| 598 |
+
# Process action embeddings
|
| 599 |
+
if noisy_actions is not None:
|
| 600 |
+
# Get mask corresponding to all action tokens
|
| 601 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 602 |
+
|
| 603 |
+
# Reshape noisy actions into individual action tokens
|
| 604 |
+
# noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
|
| 605 |
+
B = noisy_actions.shape[0]
|
| 606 |
+
noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 607 |
+
|
| 608 |
+
# Project noisy action tokens into language model embedding space
|
| 609 |
+
noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim)
|
| 610 |
+
|
| 611 |
+
# Replace embeddings of the action tokens with noisy action embeddings
|
| 612 |
+
input_embeddings = self._replace_input_embeddings(
|
| 613 |
+
input_embeddings, all_actions_mask, noisy_action_features
|
| 614 |
+
)
|
| 615 |
+
else:
|
| 616 |
+
# Replace the embeddings of the action tokens with zeros
|
| 617 |
+
# (Later on, the positional embeddings will be added to them)
|
| 618 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 619 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 620 |
+
|
| 621 |
+
# Build multimodal embeddings & attention mask
|
| 622 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 623 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 624 |
+
)
|
| 625 |
+
|
| 626 |
+
# Build labels for multimodal sequence if needed
|
| 627 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 628 |
+
|
| 629 |
+
# Dispatch to language model
|
| 630 |
+
language_model_output = self.language_model(
|
| 631 |
+
input_ids=None,
|
| 632 |
+
attention_mask=multimodal_attention_mask,
|
| 633 |
+
position_ids=None,
|
| 634 |
+
past_key_values=None,
|
| 635 |
+
inputs_embeds=multimodal_embeddings,
|
| 636 |
+
labels=multimodal_labels,
|
| 637 |
+
use_cache=use_cache,
|
| 638 |
+
output_attentions=output_attentions,
|
| 639 |
+
output_hidden_states=output_hidden_states,
|
| 640 |
+
return_dict=return_dict,
|
| 641 |
+
)
|
| 642 |
+
|
| 643 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 644 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 645 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 646 |
+
|
| 647 |
+
else:
|
| 648 |
+
raise ValueError(
|
| 649 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 650 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 651 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 652 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 653 |
+
f"=> `labels` = {labels is not None}\n"
|
| 654 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 655 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 656 |
+
f"=> `use_cache` = {use_cache}"
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 660 |
+
if not return_dict:
|
| 661 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 662 |
+
return *language_model_output, projected_patch_embeddings
|
| 663 |
+
|
| 664 |
+
return language_model_output
|
| 665 |
+
|
| 666 |
+
return PrismaticCausalLMOutputWithPast(
|
| 667 |
+
loss=language_model_output.loss,
|
| 668 |
+
logits=language_model_output.logits,
|
| 669 |
+
past_key_values=language_model_output.past_key_values,
|
| 670 |
+
hidden_states=language_model_output.hidden_states,
|
| 671 |
+
attentions=language_model_output.attentions,
|
| 672 |
+
projector_features=projected_patch_embeddings,
|
| 673 |
+
)
|
| 674 |
+
|
| 675 |
+
# === GenerationMixin Methods ===
|
| 676 |
+
def prepare_inputs_for_generation(
|
| 677 |
+
self,
|
| 678 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 679 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 680 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 681 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 682 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 683 |
+
**kwargs: str,
|
| 684 |
+
) -> Dict[str, torch.Tensor]:
|
| 685 |
+
"""Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
|
| 686 |
+
if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
|
| 687 |
+
(inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
|
| 688 |
+
):
|
| 689 |
+
raise ValueError("Generation with batch size > 1 is not currently supported!")
|
| 690 |
+
|
| 691 |
+
# Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
|
| 692 |
+
if past_key_values is not None:
|
| 693 |
+
input_ids = input_ids[:, -1:]
|
| 694 |
+
|
| 695 |
+
# If `input_embeds` are passed, we only want to use them in the 1st generation step
|
| 696 |
+
if inputs_embeds is not None and past_key_values is None:
|
| 697 |
+
model_inputs = {"input_embeds": inputs_embeds}
|
| 698 |
+
else:
|
| 699 |
+
model_inputs = {"input_ids": input_ids}
|
| 700 |
+
|
| 701 |
+
# Make sure `pixel_values` are preserved in `model_inputs`
|
| 702 |
+
model_inputs.update(
|
| 703 |
+
{
|
| 704 |
+
"attention_mask": attention_mask,
|
| 705 |
+
"pixel_values": pixel_values,
|
| 706 |
+
"past_key_values": past_key_values,
|
| 707 |
+
"use_cache": kwargs.get("use_cache"),
|
| 708 |
+
}
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
return model_inputs
|
| 712 |
+
|
| 713 |
+
# Defer to Language Model (all handle this differently, with different return types)
|
| 714 |
+
def _reorder_cache(self, *args, **kwargs) -> Any:
|
| 715 |
+
return self.language_model._reorder_cache(*args, **kwargs)
|
| 716 |
+
|
| 717 |
+
|
| 718 |
+
class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
|
| 719 |
+
config_class: PretrainedConfig = OpenVLAConfig
|
| 720 |
+
|
| 721 |
+
def __init__(self, config: OpenVLAConfig) -> None:
|
| 722 |
+
super().__init__(config)
|
| 723 |
+
self.norm_stats = config.norm_stats
|
| 724 |
+
|
| 725 |
+
# Compute action bins
|
| 726 |
+
self.bins = np.linspace(-1, 1, config.n_action_bins)
|
| 727 |
+
self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
|
| 728 |
+
|
| 729 |
+
# Compute vocab size for de-tokenization -- revert added "multiple of"
|
| 730 |
+
self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
|
| 731 |
+
|
| 732 |
+
def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
|
| 733 |
+
# eval 会用到这里
|
| 734 |
+
"""Prepares input for action prediction by adding necessary tokens"""
|
| 735 |
+
# Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
|
| 736 |
+
placeholder_action_token_ids = (
|
| 737 |
+
torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
|
| 738 |
+
)
|
| 739 |
+
input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) # torch.Size([1, 35 + 56= 91])
|
| 740 |
+
|
| 741 |
+
# Extend the attention mask to fit the new shape of input
|
| 742 |
+
# Note: Only batch size == 1 supported right now
|
| 743 |
+
mask_extension = (
|
| 744 |
+
torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
|
| 745 |
+
.to(attention_mask.device)
|
| 746 |
+
.to(attention_mask.dtype)
|
| 747 |
+
)
|
| 748 |
+
attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
|
| 749 |
+
|
| 750 |
+
return input_ids, attention_mask
|
| 751 |
+
|
| 752 |
+
def _prepare_labels_for_action_prediction(self, labels, input_ids):
|
| 753 |
+
"""Creates labels tensor for action prediction if not provided"""
|
| 754 |
+
# eval 会用到这里 ,
|
| 755 |
+
# Extends label tensors with fake action labels
|
| 756 |
+
# Adds stop tokens at the end of sequences
|
| 757 |
+
# Handles label preparation for action prediction tasks
|
| 758 |
+
# 他为啥可以随便一个? xuan说 你自定义一个值 ,然后一直指定这个 , PAD token可以吗?
|
| 759 |
+
#TODO: 这里是否要改? 感觉不需要改. 随便写就行了因为labels不重要只是要一个mask. 为什么需要这个函数? 确保 action 预测任务的标签（labels）符合模型的输入长度，并正确地处理序列终止
|
| 760 |
+
# Extend labels tensor with fake action labels
|
| 761 |
+
ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_IDX # = 为了mask正确生成, action_tokens_only_mask = (labels == ACTION_TOKEN_IDX ), 所以这里也填上ACTION_TOKEN_IDX
|
| 762 |
+
labels_extension = (
|
| 763 |
+
torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
|
| 764 |
+
* ARBITRARY_ACTION_TOKEN_IDX
|
| 765 |
+
) #torch.Size([1, 57]),全是 ARBITRARY_ACTION_TOKEN_IDX
|
| 766 |
+
labels = torch.cat([labels, labels_extension], dim=-1)
|
| 767 |
+
|
| 768 |
+
return labels
|
| 769 |
+
|
| 770 |
+
def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
|
| 771 |
+
"""Unnormalize actions using dataset statistics"""
|
| 772 |
+
action_norm_stats = self.get_action_stats(unnorm_key)
|
| 773 |
+
|
| 774 |
+
if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
|
| 775 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
|
| 776 |
+
action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
|
| 777 |
+
elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
|
| 778 |
+
mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
|
| 779 |
+
action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
|
| 780 |
+
else:
|
| 781 |
+
raise ValueError("Unsupported action/proprio normalization type detected!")
|
| 782 |
+
|
| 783 |
+
actions = np.where(
|
| 784 |
+
mask,
|
| 785 |
+
0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
|
| 786 |
+
normalized_actions,
|
| 787 |
+
)
|
| 788 |
+
|
| 789 |
+
return actions
|
| 790 |
+
|
| 791 |
+
def _run_diffusion_prediction(
|
| 792 |
+
self,
|
| 793 |
+
input_embeddings,
|
| 794 |
+
all_actions_mask,
|
| 795 |
+
noise,
|
| 796 |
+
action_head,
|
| 797 |
+
projected_patch_embeddings,
|
| 798 |
+
labels,
|
| 799 |
+
attention_mask,
|
| 800 |
+
NUM_PATCHES,
|
| 801 |
+
NUM_PROMPT_TOKENS,
|
| 802 |
+
noisy_action_projector,
|
| 803 |
+
):
|
| 804 |
+
"""Run diffusion-based action prediction"""
|
| 805 |
+
# Set diffusion timestep values
|
| 806 |
+
action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)
|
| 807 |
+
# Clone embedding for reuse in each timestep
|
| 808 |
+
orig_projected_patch_embeddings = projected_patch_embeddings.clone()
|
| 809 |
+
curr_noisy_actions = noise
|
| 810 |
+
|
| 811 |
+
# Reverse diffusion: Iteratively denoise to generate action prediction
|
| 812 |
+
for t in action_head.noise_scheduler.timesteps:
|
| 813 |
+
# Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
|
| 814 |
+
# embedding, and diffusion timestep embedding)
|
| 815 |
+
timesteps = torch.Tensor([t]).to(labels.device)
|
| 816 |
+
diffusion_timestep_embeddings = (
|
| 817 |
+
action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
|
| 818 |
+
) # (B, llm_dim)
|
| 819 |
+
diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim)
|
| 820 |
+
|
| 821 |
+
# [Diffusion] Replace the embeddings of the action tokens with noisy actions
|
| 822 |
+
# (Later on, the positional embeddings will be added to them)
|
| 823 |
+
|
| 824 |
+
# For simplicity, append diffusion timestep embedding to the end of projected vision tokens
|
| 825 |
+
projected_patch_embeddings = torch.cat(
|
| 826 |
+
(orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
|
| 827 |
+
)
|
| 828 |
+
|
| 829 |
+
# Reshape and project noisy actions into language embedding space
|
| 830 |
+
B = curr_noisy_actions.shape[0]
|
| 831 |
+
orig_curr_noisy_actions_shape = curr_noisy_actions.shape
|
| 832 |
+
curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
|
| 833 |
+
noisy_action_features = noisy_action_projector(curr_noisy_actions)
|
| 834 |
+
curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
|
| 835 |
+
|
| 836 |
+
# Replace action token embeddings with noisy action embeddings
|
| 837 |
+
input_embeddings = self._replace_input_embeddings(
|
| 838 |
+
input_embeddings.clone(), all_actions_mask, noisy_action_features
|
| 839 |
+
)
|
| 840 |
+
|
| 841 |
+
# Build multimodal embeddings and attention mask
|
| 842 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 843 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 844 |
+
)
|
| 845 |
+
|
| 846 |
+
# Forward pass through language model
|
| 847 |
+
language_model_output = self.language_model(
|
| 848 |
+
input_ids=None,
|
| 849 |
+
attention_mask=multimodal_attention_mask,
|
| 850 |
+
position_ids=None,
|
| 851 |
+
past_key_values=None,
|
| 852 |
+
inputs_embeds=multimodal_embeddings,
|
| 853 |
+
labels=None,
|
| 854 |
+
use_cache=None,
|
| 855 |
+
output_attentions=False,
|
| 856 |
+
output_hidden_states=True,
|
| 857 |
+
return_dict=True,
|
| 858 |
+
)
|
| 859 |
+
|
| 860 |
+
# Extract hidden states for action portion of response
|
| 861 |
+
last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 862 |
+
actions_hidden_states = last_hidden_states[
|
| 863 |
+
:,
|
| 864 |
+
NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
|
| 865 |
+
:,
|
| 866 |
+
] # (B, act_chunk_len, D)
|
| 867 |
+
|
| 868 |
+
# Predict noise and update noisy actions: x_t -> x_{t-1}
|
| 869 |
+
noise_pred = action_head.predict_noise(actions_hidden_states)
|
| 870 |
+
curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
|
| 871 |
+
|
| 872 |
+
curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 873 |
+
|
| 874 |
+
# Return final actions
|
| 875 |
+
return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
|
| 876 |
+
|
| 877 |
+
def _regression_or_discrete_prediction(
|
| 878 |
+
self,
|
| 879 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 880 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 881 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 882 |
+
attention_mask: torch.BoolTensor,
|
| 883 |
+
labels: torch.LongTensor,
|
| 884 |
+
NUM_PATCHES: int,
|
| 885 |
+
NUM_PROMPT_TOKENS: int,
|
| 886 |
+
action_head: L1RegressionActionHead,
|
| 887 |
+
):
|
| 888 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 889 |
+
# Extract hidden states for action tokens
|
| 890 |
+
# last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D)
|
| 891 |
+
|
| 892 |
+
# from transformers import AutoProcessor
|
| 893 |
+
# processor = AutoProcessor.from_pretrained("/data/juyi/openvla-7b+fractal20220817_data+b32+lr-5e-05+lora-r32+dropout-0.0--image_aug--test")
|
| 894 |
+
# tokenizer=processor.tokenizer
|
| 895 |
+
# tokenizer.decode(language_model_output.sequences[0])
|
| 896 |
+
|
| 897 |
+
# actions_hidden_states = last_hidden_states[:, NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + NUM_ACTIONS_CHUNK * tokennum, :]# (B, act_chunk_len, D)
|
| 898 |
+
# 都不需要取了, 直接就给 token对应的hidden state了 ,太方便了.
|
| 899 |
+
# 为什么第一个是torch.Size([1, 535, 4096])? 我应该选哪个? https://discuss.huggingface.co/t/get-each-generated-token-last-layer-hidden-state/145921
|
| 900 |
+
# language_model_output.sequences tensor([[29871, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 32001, 2]], device='cuda:0')
|
| 901 |
+
|
| 902 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 903 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 904 |
+
)
|
| 905 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 906 |
+
if self.preset:
|
| 907 |
+
# start_prefill = torch.cuda.Event(enable_timing=True)
|
| 908 |
+
# end_prefill = torch.cuda.Event(enable_timing=True)
|
| 909 |
+
# start_prefill.record()
|
| 910 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 911 |
+
# is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 912 |
+
hidden_states = language_model_output.hidden_states[0][-1]
|
| 913 |
+
actions_hidden_states = hidden_states[:, -NUM_ACTIONS_CHUNK:]
|
| 914 |
+
# end_prefill.record()
|
| 915 |
+
# torch.cuda.synchronize()
|
| 916 |
+
# prefill_time = start_prefill.elapsed_time(end_prefill) / 1000
|
| 917 |
+
# print(f"Prefill time: {prefill_time:.4f} seconds")
|
| 918 |
+
else:
|
| 919 |
+
# start_generate = torch.cuda.Event(enable_timing=True)
|
| 920 |
+
# end_generate = torch.cuda.Event(enable_timing=True)
|
| 921 |
+
# start_generate.record()
|
| 922 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=2048,output_hidden_states=True,return_dict_in_generate=True,use_cache=True)
|
| 923 |
+
# end_generate.record()
|
| 924 |
+
# torch.cuda.synchronize()
|
| 925 |
+
# generate_time = start_generate.elapsed_time(end_generate) / 1000
|
| 926 |
+
# print(f"prefill + Generate time: {generate_time:.4f} seconds")
|
| 927 |
+
actions_hidden_states = torch.stack([language_model_output.hidden_states[i][-1] for i in range(1,NUM_ACTIONS_CHUNK+1)], dim=0) # (action_chunk, batch size, seqence length, hidden_dim)
|
| 928 |
+
actions_hidden_states = actions_hidden_states.transpose(0, 1).squeeze(2) #torch.Size([batch size, action_chunk, hidden_dim])
|
| 929 |
+
|
| 930 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 931 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 932 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 933 |
+
|
| 934 |
+
return normalized_actions, actions_hidden_states
|
| 935 |
+
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def mul_regression_or_discrete_prediction(
|
| 939 |
+
self,
|
| 940 |
+
input_embeddings: torch.FloatTensor, #lanage instruction 的embedding.
|
| 941 |
+
all_actions_mask : Optional[torch.BoolTensor], #有啥用? 就是为了提取前面的embedding用. 去掉action .
|
| 942 |
+
projected_patch_embeddings: torch.FloatTensor,
|
| 943 |
+
attention_mask: torch.BoolTensor,
|
| 944 |
+
labels: torch.LongTensor,
|
| 945 |
+
NUM_PATCHES: int,
|
| 946 |
+
NUM_PROMPT_TOKENS: int,
|
| 947 |
+
action_head: L1RegressionActionHead,
|
| 948 |
+
):
|
| 949 |
+
"""Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
|
| 950 |
+
|
| 951 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 952 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 953 |
+
)
|
| 954 |
+
# multimodal_embeddings 例子'<s> <512 image token> <pripor token> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 955 |
+
language_model_output = self.language_model.generate(inputs_embeds=multimodal_embeddings,max_new_tokens=1,output_hidden_states=True,return_dict_in_generate=True)
|
| 956 |
+
# 0 , is tuple (1 token, 33 layers, torch.Size([1, 314, 4096]))
|
| 957 |
+
# import ipdb; ipdb.set_trace()
|
| 958 |
+
actions_hidden_states = language_model_output.hidden_states[0][-1]
|
| 959 |
+
actions_hidden_states = actions_hidden_states[:, -1]
|
| 960 |
+
|
| 961 |
+
normalized_actions = action_head.predict_action(actions_hidden_states)
|
| 962 |
+
normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
|
| 963 |
+
normalized_actions = normalized_actions.float().cpu().detach().numpy()
|
| 964 |
+
|
| 965 |
+
return normalized_actions, actions_hidden_states
|
| 966 |
+
|
| 967 |
+
def predict_action(
|
| 968 |
+
self,
|
| 969 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 970 |
+
unnorm_key: Optional[str] = None,
|
| 971 |
+
proprio=None,
|
| 972 |
+
proprio_projector=None,
|
| 973 |
+
action_head=None,
|
| 974 |
+
noisy_action_projector=None,
|
| 975 |
+
use_film: bool = False,
|
| 976 |
+
**kwargs: str,
|
| 977 |
+
) -> np.ndarray:
|
| 978 |
+
"""Predict actions from input sequence, with options for different prediction methods.
|
| 979 |
+
|
| 980 |
+
Args:
|
| 981 |
+
input_ids: Input token ids
|
| 982 |
+
unnorm_key: Key for unnormalization statistics
|
| 983 |
+
proprio: Proprioceptive features
|
| 984 |
+
proprio_projector: Projector for proprioceptive features
|
| 985 |
+
action_head: Optional head for L1 regression or diffusion-based prediction
|
| 986 |
+
noisy_action_projector: Projector for noisy actions in diffusion-based prediction
|
| 987 |
+
use_film: Whether to use FiLM conditioning
|
| 988 |
+
**kwargs: Additional arguments including pixel_values and attention_mask
|
| 989 |
+
|
| 990 |
+
Returns:
|
| 991 |
+
Tuple of (unnormalized_actions, action_hidden_states)
|
| 992 |
+
"""
|
| 993 |
+
# If the special empty token ('') does not already appear after the colon (':') token in the prompt
|
| 994 |
+
# (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
|
| 995 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 996 |
+
input_ids = torch.cat(
|
| 997 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 998 |
+
)
|
| 999 |
+
|
| 1000 |
+
pixel_values = kwargs["pixel_values"]
|
| 1001 |
+
attention_mask = kwargs["attention_mask"]
|
| 1002 |
+
|
| 1003 |
+
# Create fake labels tensor (needed for action mask)
|
| 1004 |
+
labels = input_ids.clone()
|
| 1005 |
+
labels[:] = IGNORE_INDEX # 输入都ignore IGNORE_INDEX = -100
|
| 1006 |
+
|
| 1007 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1008 |
+
NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1009 |
+
|
| 1010 |
+
# Prepare inputs by adding necessary tokens
|
| 1011 |
+
input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
|
| 1012 |
+
|
| 1013 |
+
# Update labels tensor for action mask computation later
|
| 1014 |
+
labels = self._prepare_labels_for_action_prediction(labels, input_ids)
|
| 1015 |
+
|
| 1016 |
+
# Get input embeddings and action masks
|
| 1017 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1018 |
+
all_actions_mask = self._process_action_masks(labels)
|
| 1019 |
+
|
| 1020 |
+
# Extract language embeddings
|
| 1021 |
+
language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1022 |
+
input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1023 |
+
)
|
| 1024 |
+
|
| 1025 |
+
# Process vision features
|
| 1026 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
|
| 1027 |
+
|
| 1028 |
+
# Add proprioceptive features if provided
|
| 1029 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1030 |
+
if use_proprio:
|
| 1031 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1032 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1033 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1034 |
+
)
|
| 1035 |
+
|
| 1036 |
+
# Use diffusion if provided, otherwise use regression or discrete prediction
|
| 1037 |
+
use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
|
| 1038 |
+
|
| 1039 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1040 |
+
NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1041 |
+
if use_proprio:
|
| 1042 |
+
NUM_PATCHES += 1
|
| 1043 |
+
|
| 1044 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1045 |
+
input_embeddings,
|
| 1046 |
+
all_actions_mask,
|
| 1047 |
+
projected_patch_embeddings,
|
| 1048 |
+
attention_mask,
|
| 1049 |
+
labels,
|
| 1050 |
+
NUM_PATCHES,
|
| 1051 |
+
NUM_PROMPT_TOKENS,
|
| 1052 |
+
action_head,
|
| 1053 |
+
)
|
| 1054 |
+
|
| 1055 |
+
# Unnormalize predicted actions
|
| 1056 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1057 |
+
|
| 1058 |
+
return actions, actions_hidden_states
|
| 1059 |
+
|
| 1060 |
+
@staticmethod
|
| 1061 |
+
def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
|
| 1062 |
+
"""Validate and resolve the unnormalization key for action statistics"""
|
| 1063 |
+
if unnorm_key is None:
|
| 1064 |
+
assert len(norm_stats) == 1, (
|
| 1065 |
+
f"Your model was trained on more than one dataset, "
|
| 1066 |
+
f"please pass a `unnorm_key` from the following options to choose the statistics "
|
| 1067 |
+
f"used for un-normalizing actions: {norm_stats.keys()}"
|
| 1068 |
+
)
|
| 1069 |
+
unnorm_key = next(iter(norm_stats.keys()))
|
| 1070 |
+
# norm states没有加载libero, 为什么?
|
| 1071 |
+
assert unnorm_key in norm_stats, (
|
| 1072 |
+
f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
|
| 1073 |
+
f"please choose from: {norm_stats.keys()}"
|
| 1074 |
+
)
|
| 1075 |
+
return unnorm_key
|
| 1076 |
+
|
| 1077 |
+
def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
|
| 1078 |
+
"""Get the dimensionality of the policy's action space."""
|
| 1079 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1080 |
+
return len(self.norm_stats[unnorm_key]["action"]["min"])
|
| 1081 |
+
|
| 1082 |
+
def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
|
| 1083 |
+
"""Get all the logged statistics for the given dataset."""
|
| 1084 |
+
unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
|
| 1085 |
+
return self.norm_stats[unnorm_key]["action"]
|
| 1086 |
+
|
| 1087 |
+
|
| 1088 |
+
def lisa_forward(
|
| 1089 |
+
self,
|
| 1090 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1091 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1092 |
+
pixel_values: Optional[torch.FloatTensor] = None,
|
| 1093 |
+
labels: Optional[torch.LongTensor] = None,
|
| 1094 |
+
inputs_embeds: Optional[torch.FloatTensor] = None,
|
| 1095 |
+
past_key_values: Optional[List[torch.FloatTensor]] = None,
|
| 1096 |
+
use_cache: Optional[bool] = None,
|
| 1097 |
+
output_attentions: Optional[bool] = None,
|
| 1098 |
+
output_hidden_states: Optional[bool] = None,
|
| 1099 |
+
output_projector_features: Optional[bool] = None,
|
| 1100 |
+
return_dict: Optional[bool] = None,
|
| 1101 |
+
proprio=None,
|
| 1102 |
+
proprio_projector=None,
|
| 1103 |
+
**kwargs
|
| 1104 |
+
) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
|
| 1105 |
+
"""Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
|
| 1106 |
+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
| 1107 |
+
output_hidden_states = (
|
| 1108 |
+
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
| 1109 |
+
)
|
| 1110 |
+
output_projector_features = output_projector_features if output_projector_features is not None else False
|
| 1111 |
+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
| 1112 |
+
|
| 1113 |
+
# Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
|
| 1114 |
+
use_cache = use_cache and not self.training
|
| 1115 |
+
|
| 1116 |
+
# Instantiate Placeholder for Projector Features
|
| 1117 |
+
projected_patch_embeddings = None
|
| 1118 |
+
|
| 1119 |
+
# === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
|
| 1120 |
+
if input_ids.shape[1] == 1:
|
| 1121 |
+
assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
|
| 1122 |
+
assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
|
| 1123 |
+
assert labels is None, "Unexpected key `labels` provided during cached generation!"
|
| 1124 |
+
|
| 1125 |
+
language_model_output = self.language_model(
|
| 1126 |
+
input_ids=input_ids,
|
| 1127 |
+
attention_mask=None,
|
| 1128 |
+
position_ids=None,
|
| 1129 |
+
past_key_values=past_key_values,
|
| 1130 |
+
inputs_embeds=None,
|
| 1131 |
+
labels=None,
|
| 1132 |
+
use_cache=use_cache,
|
| 1133 |
+
output_attentions=output_attentions,
|
| 1134 |
+
output_hidden_states=output_hidden_states,
|
| 1135 |
+
return_dict=return_dict,
|
| 1136 |
+
)
|
| 1137 |
+
|
| 1138 |
+
# === Handle Unimodal Forward ===
|
| 1139 |
+
elif pixel_values is None:
|
| 1140 |
+
raise NotImplementedError
|
| 1141 |
+
|
| 1142 |
+
# === Handle Multimodal Forward ===
|
| 1143 |
+
elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
|
| 1144 |
+
assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
|
| 1145 |
+
|
| 1146 |
+
# Get input embeddings (from language model embeddings)
|
| 1147 |
+
input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D)
|
| 1148 |
+
# Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
|
| 1149 |
+
# language_embeddings = input_embeddings[~all_actions_mask].reshape(
|
| 1150 |
+
# input_embeddings.shape[0], -1, input_embeddings.shape[2]
|
| 1151 |
+
# ) # (B, lang_seq_len, llm_dim) 这里就会把结尾的 stop index和padding 也算进去. 没问题吗? 没问题因为ignore了 我直接删了因为不用film
|
| 1152 |
+
# Get visual features
|
| 1153 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1154 |
+
|
| 1155 |
+
# Add proprioceptive state if provided
|
| 1156 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1157 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1158 |
+
)
|
| 1159 |
+
|
| 1160 |
+
all_actions_mask = (labels == ACTION_TOKEN_IDX) #和run forward pass不一样, run forward pass要手动算token number来找偏移.
|
| 1161 |
+
all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1)
|
| 1162 |
+
input_embeddings = input_embeddings * ~all_actions_mask
|
| 1163 |
+
|
| 1164 |
+
# Build multimodal embeddings & attention mask
|
| 1165 |
+
multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
|
| 1166 |
+
input_embeddings, projected_patch_embeddings, attention_mask
|
| 1167 |
+
)
|
| 1168 |
+
|
| 1169 |
+
# Build labels for multimodal sequence if needed
|
| 1170 |
+
multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
|
| 1171 |
+
|
| 1172 |
+
# Dispatch to language model
|
| 1173 |
+
language_model_output = self.language_model(
|
| 1174 |
+
input_ids=None,
|
| 1175 |
+
attention_mask=multimodal_attention_mask,
|
| 1176 |
+
position_ids=None,
|
| 1177 |
+
past_key_values=None,
|
| 1178 |
+
inputs_embeds=multimodal_embeddings,
|
| 1179 |
+
labels=multimodal_labels,
|
| 1180 |
+
use_cache=use_cache,
|
| 1181 |
+
output_attentions=output_attentions,
|
| 1182 |
+
output_hidden_states=output_hidden_states,
|
| 1183 |
+
return_dict=return_dict,
|
| 1184 |
+
)
|
| 1185 |
+
|
| 1186 |
+
|
| 1187 |
+
# === Otherwise =>> Assume Invalid! ===
|
| 1188 |
+
elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
|
| 1189 |
+
raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
|
| 1190 |
+
|
| 1191 |
+
else:
|
| 1192 |
+
raise ValueError(
|
| 1193 |
+
"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
|
| 1194 |
+
f"=> `input_ids` = {input_ids is not None}\n"
|
| 1195 |
+
f"=> `attention_mask` = {attention_mask is not None}\n"
|
| 1196 |
+
f"=> `pixel_values` = {pixel_values is not None}\n"
|
| 1197 |
+
f"=> `labels` = {labels is not None}\n"
|
| 1198 |
+
f"=> `input_embeds` = {inputs_embeds is not None}\n"
|
| 1199 |
+
f"=> `past_key_values` = {past_key_values is not None}\n"
|
| 1200 |
+
f"=> `use_cache` = {use_cache}"
|
| 1201 |
+
)
|
| 1202 |
+
|
| 1203 |
+
# Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
|
| 1204 |
+
if not return_dict:
|
| 1205 |
+
if output_projector_features and (projected_patch_embeddings is not None):
|
| 1206 |
+
return *language_model_output, projected_patch_embeddings
|
| 1207 |
+
|
| 1208 |
+
return language_model_output
|
| 1209 |
+
|
| 1210 |
+
return PrismaticCausalLMOutputWithPast(
|
| 1211 |
+
loss=language_model_output.loss,
|
| 1212 |
+
logits=language_model_output.logits,
|
| 1213 |
+
past_key_values=language_model_output.past_key_values,
|
| 1214 |
+
hidden_states=language_model_output.hidden_states,
|
| 1215 |
+
attentions=language_model_output.attentions,
|
| 1216 |
+
projector_features=projected_patch_embeddings,
|
| 1217 |
+
)
|
| 1218 |
+
|
| 1219 |
+
def lisa_predict_action(
|
| 1220 |
+
self,
|
| 1221 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1222 |
+
unnorm_key: Optional[str] = None,
|
| 1223 |
+
proprio=None,
|
| 1224 |
+
proprio_projector=None,
|
| 1225 |
+
action_head:L1RegressionActionHead=None,
|
| 1226 |
+
noisy_action_projector=None,
|
| 1227 |
+
use_film: bool = False,
|
| 1228 |
+
**kwargs: str,
|
| 1229 |
+
) -> np.ndarray:
|
| 1230 |
+
|
| 1231 |
+
pixel_values = kwargs["pixel_values"]
|
| 1232 |
+
attention_mask = kwargs["attention_mask"]
|
| 1233 |
+
|
| 1234 |
+
# Get number of tokens in prompt (excluding the start token)
|
| 1235 |
+
# NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token
|
| 1236 |
+
|
| 1237 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1238 |
+
#预测的时候labels 有啥用? 只是用来设置mask 我们自回归就不用
|
| 1239 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1240 |
+
if cfg.preset:
|
| 1241 |
+
special_tensor = torch.tensor([[29871]], device=input_ids.device, dtype=input_ids.dtype)
|
| 1242 |
+
output_tensor = torch.tensor([[32001] * NUM_ACTIONS_CHUNK], device=input_ids.device, dtype=input_ids.dtype)
|
| 1243 |
+
input_ids = torch.cat([input_ids, special_tensor, output_tensor], dim=1) # preset action tokens, only forward once.
|
| 1244 |
+
self.preset = True
|
| 1245 |
+
else:
|
| 1246 |
+
self.preset = False
|
| 1247 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1248 |
+
|
| 1249 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1250 |
+
|
| 1251 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1252 |
+
if use_proprio:
|
| 1253 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1254 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1255 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1256 |
+
)
|
| 1257 |
+
|
| 1258 |
+
# Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
|
| 1259 |
+
# NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
|
| 1260 |
+
# if use_proprio:
|
| 1261 |
+
# NUM_PATCHES += 1
|
| 1262 |
+
normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
|
| 1263 |
+
input_embeddings,
|
| 1264 |
+
None,
|
| 1265 |
+
projected_patch_embeddings,
|
| 1266 |
+
attention_mask,
|
| 1267 |
+
None, #推理不需要labels
|
| 1268 |
+
None, #推理不需要NUM_PATCHES
|
| 1269 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1270 |
+
action_head,
|
| 1271 |
+
)
|
| 1272 |
+
|
| 1273 |
+
# Unnormalize predicted actions
|
| 1274 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1275 |
+
|
| 1276 |
+
return actions, actions_hidden_states
|
| 1277 |
+
|
| 1278 |
+
|
| 1279 |
+
def mul_predict_action(
|
| 1280 |
+
self,
|
| 1281 |
+
input_ids: Optional[torch.LongTensor] = None, #就是 language instruction.
|
| 1282 |
+
unnorm_key: Optional[str] = None,
|
| 1283 |
+
proprio=None,
|
| 1284 |
+
proprio_projector=None,
|
| 1285 |
+
action_head:L1RegressionActionHead=None,
|
| 1286 |
+
noisy_action_projector=None,
|
| 1287 |
+
use_film: bool = False,
|
| 1288 |
+
**kwargs: str,
|
| 1289 |
+
) -> np.ndarray:
|
| 1290 |
+
|
| 1291 |
+
if not torch.all(input_ids[:, -1] == 29871):
|
| 1292 |
+
input_ids = torch.cat(
|
| 1293 |
+
(input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
|
| 1294 |
+
)
|
| 1295 |
+
|
| 1296 |
+
pixel_values = kwargs["pixel_values"]
|
| 1297 |
+
attention_mask = kwargs["attention_mask"]
|
| 1298 |
+
|
| 1299 |
+
# input id '<s> In: What action should the robot take to open the middle drawer of the cabinet?\nOut:'
|
| 1300 |
+
cfg = kwargs.get("cfg", None) # Extract cfg from kwargs
|
| 1301 |
+
input_embeddings = self.get_input_embeddings()(input_ids)
|
| 1302 |
+
|
| 1303 |
+
projected_patch_embeddings = self._process_vision_features(pixel_values)
|
| 1304 |
+
|
| 1305 |
+
use_proprio = proprio_projector is not None and proprio is not None
|
| 1306 |
+
if use_proprio:
|
| 1307 |
+
proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
|
| 1308 |
+
projected_patch_embeddings = self._process_proprio_features(
|
| 1309 |
+
projected_patch_embeddings, proprio, proprio_projector
|
| 1310 |
+
)
|
| 1311 |
+
|
| 1312 |
+
normalized_actions, actions_hidden_states = self.mul_regression_or_discrete_prediction(
|
| 1313 |
+
input_embeddings,
|
| 1314 |
+
None,
|
| 1315 |
+
projected_patch_embeddings,
|
| 1316 |
+
attention_mask,
|
| 1317 |
+
None, #推理不需要labels
|
| 1318 |
+
None, #推理不需要NUM_PATCHES
|
| 1319 |
+
None, #推理不需要NUM_PROMPT_TOKENS
|
| 1320 |
+
action_head,
|
| 1321 |
+
)
|
| 1322 |
+
|
| 1323 |
+
# Unnormalize predicted actions
|
| 1324 |
+
actions = self._unnormalize_actions(normalized_actions, unnorm_key)
|
| 1325 |
+
|
| 1326 |
+
return actions, actions_hidden_states
|
| 1327 |
+
|
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 |
+
}
|
special_tokens_map.json
ADDED
|
@@ -0,0 +1,39 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"additional_special_tokens": [
|
| 3 |
+
{
|
| 4 |
+
"content": "<ACT>",
|
| 5 |
+
"lstrip": false,
|
| 6 |
+
"normalized": false,
|
| 7 |
+
"rstrip": false,
|
| 8 |
+
"single_word": false
|
| 9 |
+
}
|
| 10 |
+
],
|
| 11 |
+
"bos_token": {
|
| 12 |
+
"content": "<s>",
|
| 13 |
+
"lstrip": false,
|
| 14 |
+
"normalized": false,
|
| 15 |
+
"rstrip": false,
|
| 16 |
+
"single_word": false
|
| 17 |
+
},
|
| 18 |
+
"eos_token": {
|
| 19 |
+
"content": "</s>",
|
| 20 |
+
"lstrip": false,
|
| 21 |
+
"normalized": false,
|
| 22 |
+
"rstrip": false,
|
| 23 |
+
"single_word": false
|
| 24 |
+
},
|
| 25 |
+
"pad_token": {
|
| 26 |
+
"content": "<PAD>",
|
| 27 |
+
"lstrip": false,
|
| 28 |
+
"normalized": false,
|
| 29 |
+
"rstrip": false,
|
| 30 |
+
"single_word": false
|
| 31 |
+
},
|
| 32 |
+
"unk_token": {
|
| 33 |
+
"content": "<unk>",
|
| 34 |
+
"lstrip": false,
|
| 35 |
+
"normalized": false,
|
| 36 |
+
"rstrip": false,
|
| 37 |
+
"single_word": false
|
| 38 |
+
}
|
| 39 |
+
}
|
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,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+
"32001": {
|
| 38 |
+
"content": "<ACT>",
|
| 39 |
+
"lstrip": false,
|
| 40 |
+
"normalized": false,
|
| 41 |
+
"rstrip": false,
|
| 42 |
+
"single_word": false,
|
| 43 |
+
"special": true
|
| 44 |
+
}
|
| 45 |
+
},
|
| 46 |
+
"additional_special_tokens": [
|
| 47 |
+
"<ACT>"
|
| 48 |
+
],
|
| 49 |
+
"auto_map": {
|
| 50 |
+
"AutoProcessor": "processing_prismatic.PrismaticProcessor"
|
| 51 |
+
},
|
| 52 |
+
"bos_token": "<s>",
|
| 53 |
+
"clean_up_tokenization_spaces": false,
|
| 54 |
+
"eos_token": "</s>",
|
| 55 |
+
"legacy": false,
|
| 56 |
+
"model_max_length": 2048,
|
| 57 |
+
"pad_token": "<PAD>",
|
| 58 |
+
"padding_side": "right",
|
| 59 |
+
"processor_class": "PrismaticProcessor",
|
| 60 |
+
"sp_model_kwargs": {},
|
| 61 |
+
"tokenizer_class": "LlamaTokenizer",
|
| 62 |
+
"unk_token": "<unk>",
|
| 63 |
+
"use_default_system_prompt": false
|
| 64 |
+
}
|