vla added but giving nans in loss

added_tokens.json

@@ -7,6 +7,8 @@
   "<|im_end|>": 151645,
   "<|im_start|>": 151644,
   "<|image|>": 152068,
+  "<|proprio|>": 152069,
+  "<|skill|>": 152070,
   "|<EXTRA_TOKENS_0>|": 151646,
   "|<EXTRA_TOKENS_100>|": 151746,
   "|<EXTRA_TOKENS_101>|": 151747,

config.json

@@ -28,5 +28,7 @@
   "use_cache": true,
   "use_position_ids": true,
   "vocab_size": 152064,
+  "skill_vocab_size": 1000,
+  "additional_vocab_size": 128,
   "weight_tying": false
 }

config_molmo.py

@@ -9,6 +9,8 @@ class MolmoConfig(PretrainedConfig):
 
     def __init__(
         self,
+        skill_vocab_size=1000,
+        additional_vocab_size=128,
         vocab_size=50304,
         embedding_size=50304,
         hidden_size=4096,
@@ -31,6 +33,8 @@ class MolmoConfig(PretrainedConfig):
         layer_norm_type: str="rms",
         **kwargs,
     ):
+        self.skill_vocab_size = skill_vocab_size
+        self.additional_vocab_size = additional_vocab_size
         self.vocab_size = vocab_size
         self.embedding_size = embedding_size
         self.max_position_embeddings = max_position_embeddings

modeling_molmo.py

@@ -541,6 +541,7 @@ class Embedding(nn.Module):
         self,
         num_embeddings: int,
         num_new_embeddings: int,
+        num_skill_embeddings: int,
         features: int,
         device: Union[str, torch.device],
         initializer_range: float = 0.02,
@@ -555,13 +556,17 @@ class Embedding(nn.Module):
         self.new_embedding = nn.Parameter(
             torch.zeros(num_new_embeddings, features, device=device),
         )
+        self.skill_embedding = nn.Parameter(
+            torch.zeros(num_skill_embeddings, features, device=device),
+        )
 
     def reset_parameters(self):
         nn.init.normal_(self.embedding, std=self.initializer_range)
         nn.init.normal_(self.new_embedding, std=self.new_embed_initializer_range)
+        nn.init.normal_(self.skill_embedding, std=self.new_embed_initializer_range)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return F.embedding(x, torch.cat([self.embedding, self.new_embedding], dim=0))
+        return F.embedding(x, torch.cat([self.embedding, self.new_embedding, self.skill_embedding], dim=0))
 
 
 class Dropout(nn.Dropout):
@@ -681,6 +686,7 @@ class FullMolmoConfig:
     initializer_range: float = 0.02
     normalize_input_embeds: bool = False
     use_position_ids: bool = True
+    skill_vocab_size: int = 1000
 
     @property
     def effective_n_kv_heads(self) -> int:
@@ -1695,6 +1701,7 @@ class Molmo(nn.Module):
             wte = Embedding(
                 config.embedding_size or config.vocab_size,
                 config.additional_vocab_size,
+                config.skill_vocab_size,
                 config.d_model,
                 device=config.init_device,
                 initializer_range=config.initializer_range,
@@ -1734,6 +1741,16 @@ class Molmo(nn.Module):
                         )
                 }
             )
+            self.transformer.update(
+                {
+                    "skill_ff_out": nn.Linear(
+                        config.d_model,
+                        config.skill_vocab_size,
+                        bias=config.include_bias,
+                        device=config.init_device,
+                        )
+                }
+            )
 
         self.vision_backbone: Optional[OLMoVisionBackbone] = None
         if config.vision_backbone is not None:
@@ -1741,6 +1758,11 @@ class Molmo(nn.Module):
 
         self.__num_fwd_flops: Optional[int] = None
 
+        self.total_vocab_size = config.vocab_size + config.additional_vocab_size + config.skill_vocab_size
+        torch.nn.init.xavier_uniform_(self.transformer.skill_ff_out.weight)
+        if self.transformer.skill_ff_out.bias is not None:
+            torch.nn.init.zeros_(self.transformer.skill_ff_out.bias)
+
     def reset_parameters(self):
         if self.vision_backbone is not None:
             self.vision_backbone.reset_parameters()
@@ -1778,12 +1800,15 @@ class Molmo(nn.Module):
         image_masks: Optional[torch.Tensor] = None,
         image_input_idx: Optional[torch.Tensor] = None,
         subsegment_ids: Optional[torch.Tensor] = None,
+        proprio_embeds: Optional[torch.Tensor] = None,
+        proprio_idx: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
         past_key_values: Optional[Sequence[Tuple[torch.Tensor, torch.Tensor]]] = None,
         use_cache: bool = False,
         last_logits_only: bool = False,
         output_hidden_states: Optional[bool] = None,
         append_last_valid_logits: Optional[torch.Tensor] = None,
+        mode: Optional[str] = "vla",
     ) -> ModelOutput:
         """
         :param input_ids: A tensor of shape `(batch_size, seq_len)`.
@@ -1880,6 +1905,9 @@ class Molmo(nn.Module):
             image_features = image_features.to(x.device)
 
             x[batch_idx[valid], image_input_idx[valid]] += image_features[valid]
+        
+        if proprio_embeds is not None:
+            x[batch_idx, proprio_idx] += proprio_embeds
 
         if not self.config.rope:
             # Get positional embeddings.
@@ -1997,7 +2025,14 @@ class Molmo(nn.Module):
         if self.config.weight_tying:
             logits = F.linear(x, self.transformer.wte.weight, None)  # type: ignore
         else:
-            logits = self.transformer.ff_out(x)  # type: ignore
+            if mode == "vla":
+                logits = self.transformer.skill_ff_out(x)
+                # this little trick allows use to use HF generate() while decoding
+                if use_cache:
+                    filler_logits = torch.full((x.shape[0], x.shape[1], self.total_vocab_size-self.config.skill_vocab_size), -math.inf, device=logits.device)
+                    logits = torch.cat([filler_logits, logits], dim=-1)  # type: ignore
+            else:
+                logits = self.transformer.ff_out(x)  # type: ignore
         if self.config.scale_logits:
             logits.mul_(1 / math.sqrt(self.config.d_model))
 
@@ -2039,6 +2074,7 @@ class MolmoForCausalLM(PreTrainedModel):
                 mlp_hidden_size=config.intermediate_size,
                 n_layers=config.num_hidden_layers,
                 additional_vocab_size=128,
+                skill_vocab_size=config.skill_vocab_size,
                 n_heads=config.num_attention_heads,
                 n_kv_heads=config.num_key_value_heads,
                 rope_theta=config.rope_theta,
@@ -2080,6 +2116,8 @@ class MolmoForCausalLM(PreTrainedModel):
         image_masks: Optional[torch.Tensor] = None,
         image_input_idx: Optional[torch.Tensor] = None,
         subsegment_ids: Optional[torch.Tensor] = None,
+        proprio_embeds: Optional[torch.Tensor] = None,
+        proprio_idx: Optional[torch.Tensor] = None,
         position_ids: Optional[torch.Tensor] = None,
         past_key_values: Optional[List[torch.FloatTensor]] = None,
         labels: Optional[torch.LongTensor] = None,
@@ -2113,6 +2151,8 @@ class MolmoForCausalLM(PreTrainedModel):
             image_masks=image_masks,
             image_input_idx=image_input_idx,
             subsegment_ids=subsegment_ids,
+            proprio_embeds=proprio_embeds,
+            proprio_idx=proprio_idx,
             position_ids=position_ids,
             past_key_values=past_key_values,
             use_cache=use_cache,
@@ -2185,6 +2225,8 @@ class MolmoForCausalLM(PreTrainedModel):
         images = batch.get("images")
         image_masks = batch.get("image_masks")
         image_input_idx = batch.get("image_input_idx")
+        proprio_embeds = batch.get("proprio_embeds")
+        proprio_idx = batch.get("proprio_idx")
 
         # Validate inputs.
         input_ids = batch["input_ids"]
@@ -2217,6 +2259,8 @@ class MolmoForCausalLM(PreTrainedModel):
             image_masks=image_masks,
             image_input_idx=image_input_idx,
             position_ids=position_ids,
+            proprio_embeds=proprio_embeds,
+            proprio_idx=proprio_idx,
             append_last_valid_logits=append_last_valid_logits,
             **kwargs,
         )
@@ -2235,6 +2279,8 @@ class MolmoForCausalLM(PreTrainedModel):
             images = kwargs.get("images")
             image_masks = kwargs.get("image_masks")
             image_input_idx = kwargs.get("image_input_idx")
+            proprio_embeds = kwargs.get("proprio_embeds")
+            proprio_idx = kwargs.get("proprio_idx")
             position_ids = kwargs.get("position_ids")
             append_last_valid_logits = kwargs.get("append_last_valid_logits")
             model_inputs = {
@@ -2250,6 +2296,8 @@ class MolmoForCausalLM(PreTrainedModel):
                 model_inputs["image_masks"] = image_masks
                 model_inputs["image_input_idx"] = image_input_idx
                 model_inputs["append_last_valid_logits"] = append_last_valid_logits
+                model_inputs["proprio_embeds"] = proprio_embeds
+                model_inputs["proprio_idx"] = proprio_idx
         else:
             model_inputs = {"input_ids": input_ids, "past_key_values": past_key_values}
 
@@ -2272,6 +2320,9 @@ class MolmoForCausalLM(PreTrainedModel):
                 del model_kwargs["images"]
                 del model_kwargs["image_masks"]
                 del model_kwargs["image_input_idx"]
+            if "proprio_embeds" in model_kwargs:
+                del model_kwargs["proprio_embeds"]
+                del model_kwargs["proprio_idx"]
         cache_name, cache = super()._extract_past_from_model_output(outputs)
         model_kwargs[cache_name] = cache
         model_kwargs["cache_position"] = model_kwargs["cache_position"][-1:] + num_new_tokens

preprocessing_molmo.py

@@ -28,7 +28,7 @@ from transformers.utils import logging
 
 from transformers import AutoTokenizer
 from .image_preprocessing_molmo import MolmoImagesKwargs, MolmoImageProcessor
-
+from typing import List, Union
 
 logger = logging.get_logger(__name__)
 
@@ -38,9 +38,14 @@ DEFAULT_IM_START_TOKEN = f"<im_start>"
 DEFAULT_IM_END_TOKEN = f"<im_end>"
 DEFAULT_IM_COL_TOKEN = f"<im_col>"
 IMAGE_PROMPT = "<|image|>"
+PROPRIO_PROMPT = "<|proprio|>"
+SKILL_PROMPT = "<|skill|>"
 
-EXTRA_TOKENS = (DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_COL_TOKEN, IMAGE_PROMPT)
+EXTRA_TOKENS = (DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_COL_TOKEN, IMAGE_PROMPT, PROPRIO_PROMPT, SKILL_PROMPT)
 
+ProprioInput = Union[
+    np.ndarray, "torch.Tensor", List[np.ndarray], List["torch.Tensor"]
+]
 
 def get_special_token_ids(tokenizer):
     ids = tokenizer.encode("".join(EXTRA_TOKENS), add_special_tokens=False)
@@ -72,7 +77,7 @@ class MolmoProcessorKwargs(ProcessingKwargs, total=False):
         "text_kwargs": {
             "style": "long_caption",
             "system_prompt": "none",
-            "message_format": "role",
+            "message_format": "robot",
             "always_start_with_space": True,
             "sequence_length": 1536,
             "padding": False,
@@ -97,11 +102,14 @@ class MolmoProcessor(ProcessorMixin):
             self._special_tokens = get_special_token_ids(self.tokenizer)
         return self._special_tokens
 
-    def get_tokens_input(self, prompt, message_format, always_start_with_space):
+    def get_tokens_input(self, prompt, message_format, always_start_with_space, num_proprio):
         if message_format == "none" or message_format is None:
             pass
         elif message_format == "role":
             prompt = "User: " + prompt + " Assistant:"
+        elif message_format == "robot":
+            # this adds proprio observations after the prompt
+            prompt = "User: " + prompt + PROPRIO_PROMPT*num_proprio + " Assistant:"
         else:
             raise NotImplementedError(f"Message format {message_format} not implemented")
 
@@ -116,6 +124,7 @@ class MolmoProcessor(ProcessorMixin):
         self,
         text: TextInput = None,
         images: ImageInput = None,
+        proprio: ProprioInput = None,
         *,
         tokens: Optional[PreTokenizedInput] = None,
         **kwargs: Unpack[MolmoProcessorKwargs],
@@ -126,14 +135,18 @@ class MolmoProcessor(ProcessorMixin):
             **kwargs,
         )
 
+        num_proprio = len(proprio) if proprio is not None else 0
+
         if tokens is None:
             tokens = self.get_tokens_input(
                 text,
                 output_kwargs["text_kwargs"]["message_format"],
                 output_kwargs["text_kwargs"]["always_start_with_space"],
+                num_proprio
             )
 
         image_token_id = self.special_token_ids[IMAGE_PROMPT]
+        proprio_token_id = self.special_token_ids[PROPRIO_PROMPT]
 
         if images is not None:
             if not isinstance(images, (list, tuple)):
@@ -182,6 +195,9 @@ class MolmoProcessor(ProcessorMixin):
             # Shift patch mapping up by one since we added BOS
             image_input_idx = out["image_input_idx"]
             out["image_input_idx"] = np.where(image_input_idx < 0, image_input_idx, image_input_idx + 1)
+        
+        proprio_idx = np.where(out["input_ids"] == proprio_token_id)[0]
+        out["proprio_idx"] = proprio_idx
 
         for k, v in out.items():
             out[k] = torch.from_numpy(v)

special_tokens_map.json

@@ -422,7 +422,9 @@
     "<im_end>",
     "<im_patch>",
     "<im_col>",
-    "<|image|>"
+    "<|image|>",
+    "<|proprio|>",
+    "<|skill|>"
   ],
   "eos_token": {
     "content": "<|endoftext|>",

tokenizer_config.json

@@ -3408,6 +3408,22 @@
       "rstrip": false,
       "single_word": false,
       "special": true
+    },
+    "152069": {
+      "content": "<|proprio|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "152070": {
+      "content": "<|skill|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
     }
   },
   "additional_special_tokens": [
@@ -3833,7 +3849,9 @@
     "<im_end>",
     "<im_patch>",
     "<im_col>",
-    "<|image|>"
+    "<|image|>",
+    "<|proprio|>",
+    "<|skill|>"
   ],
   "auto_map": {
     "AutoProcessor": "preprocessing_molmo.MolmoProcessor"