Core function for VLA inference; maps input image and task instruction to continuous action (de-tokenizes). @param image: PIL Image as [height, width, 3] @param instruction: Task instruction string @param unnorm_key: Optional dataset name for retrieving un-normalizi
(
self, image: Image, instruction: str, unnorm_key: Optional[str] = None, **kwargs: str
)
| 32 | |
| 33 | @torch.inference_mode() |
| 34 | def predict_action( |
| 35 | self, image: Image, instruction: str, unnorm_key: Optional[str] = None, **kwargs: str |
| 36 | ) -> np.ndarray: |
| 37 | """ |
| 38 | Core function for VLA inference; maps input image and task instruction to continuous action (de-tokenizes). |
| 39 | |
| 40 | @param image: PIL Image as [height, width, 3] |
| 41 | @param instruction: Task instruction string |
| 42 | @param unnorm_key: Optional dataset name for retrieving un-normalizing statistics; if None, checks that model |
| 43 | was trained only on a single dataset, and retrieves those statistics. |
| 44 | |
| 45 | @return Unnormalized (continuous) action vector --> end-effector deltas. |
| 46 | """ |
| 47 | image_transform, tokenizer = self.vision_backbone.image_transform, self.llm_backbone.tokenizer |
| 48 | |
| 49 | # Build VLA Prompt |
| 50 | prompt_builder = self.get_prompt_builder() |
| 51 | prompt_builder.add_turn(role="human", message=f"What action should the robot take to {instruction.lower()}?") |
| 52 | prompt_text = prompt_builder.get_prompt() |
| 53 | |
| 54 | # Prepare Inputs |
| 55 | input_ids = tokenizer(prompt_text, truncation=True, return_tensors="pt").input_ids.to(self.device) |
| 56 | if isinstance(tokenizer, LlamaTokenizerFast): |
| 57 | # Note: We need to add this special empty token ('') after the colon (':') token in "ASSISTANT:" |
| 58 | # in order for the predictions to match the training configuration and be accurate. |
| 59 | input_ids = torch.cat( |
| 60 | (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(self.device)), dim=1 |
| 61 | ) |
| 62 | else: |
| 63 | raise ValueError(f"Unsupported `tokenizer` type = {type(tokenizer)}") |
| 64 | |
| 65 | # Preprocess Image |
| 66 | pixel_values = image_transform(image) |
| 67 | if isinstance(pixel_values, torch.Tensor): |
| 68 | pixel_values = pixel_values[None, ...].to(self.device) |
| 69 | elif isinstance(pixel_values, dict): |
| 70 | pixel_values = {k: v[None, ...].to(self.device) for k, v in pixel_values.items()} |
| 71 | else: |
| 72 | raise ValueError(f"Unsupported `pixel_values` type = {type(pixel_values)}") |
| 73 | |
| 74 | # Invoke super().generate --> taps into `GenerationMixin` which (redirects) to `forward()` |
| 75 | autocast_dtype = self.llm_backbone.half_precision_dtype |
| 76 | with torch.autocast("cuda", dtype=autocast_dtype, enabled=self.enable_mixed_precision_training): |
| 77 | # fmt: off |
| 78 | generated_ids = super(PrismaticVLM, self).generate( |
| 79 | input_ids=input_ids, # Shape: [1, seq] |
| 80 | pixel_values=pixel_values, # Shape: [1, 3, res, res] or Dict[str, ...] |
| 81 | max_new_tokens=self.get_action_dim(unnorm_key), |
| 82 | **kwargs |
| 83 | ) |
| 84 | # fmt: on |
| 85 | |
| 86 | # Extract predicted action tokens and translate into (normalized) continuous actions |
| 87 | predicted_action_token_ids = generated_ids[0, -self.get_action_dim(unnorm_key) :] |
| 88 | normalized_actions = self.action_tokenizer.decode_token_ids_to_actions(predicted_action_token_ids.cpu().numpy()) |
| 89 | |
| 90 | # Un-normalize Actions |
| 91 | action_norm_stats = self.get_action_stats(unnorm_key) |
no test coverage detected