Update modeling_prismatic.py

modeling_prismatic.py

@@ -581,7 +581,14 @@ class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
         generated_ids = self.generate(input_ids, **kwargs)
         # VLLM inference 
         # Extract predicted action tokens and translate into (normalized) continuous actions
-        predicted_action_token_ids = generated_ids[0, -(self.get_action_dim(unnorm_key) + 1) : -1].cpu().numpy()
+        # predicted_action_token_ids = generated_ids[-1, -(self.get_action_dim(unnorm_key) + 1) : -1].cpu().numpy()
+        # Get the last row of generated_ids
+        last_row = generated_ids[-1]
+
+        # Remove padding tokens (assuming pad token id is 32000)
+        non_pad_tokens = last_row[last_row != 32000]
+        # Take the last 7 tokens from the filtered tokens
+        predicted_action_token_ids = non_pad_tokens[-(self.get_action_dim(unnorm_key) + 1) : -1].cpu().numpy()
         discretized_actions = self.vocab_size - predicted_action_token_ids
         discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
         normalized_actions = self.bin_centers[discretized_actions]