Scheduled Commit

groot/vla/model/dreamzero/transform/dreamzero_cotrain.py

@@ -126,9 +126,11 @@ def collate(features: List[dict], tokenizer: AutoTokenizer, num_views=3, embodim
                         processed_item = "A multi-view video shows that a robot " + processed_item.lower() + " The video is split into four views: The top-left view shows the top camera, the top-right view shows the right camera, the bottom-left view shows the left camera, and the bottom-right view is a black screen. The robot " + processed_item.lower()
                     elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.TREX.value]:
                         processed_item = "A multi-view video shows that a bimanual robot with dexterous hands " + processed_item.lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right wrist, the bottom-left view shows the camera view from the left wrist, and the bottom-right view is a black screen. The robot " + processed_item.lower()
+                    elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.ROBOCASA_BIMANUAL_PANDA_INSPIRE_HAND.value]:
+                        processed_item = "A multi-view video shows that a bimanual robot with dexterous hands " + processed_item.lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right wrist, the bottom-left view shows the camera view from the left wrist, and the bottom-right view is a black screen. The robot " + processed_item.lower()
                     else:
-                        raise ValueError(f"Embodiment ID {elem['embodiment_id']} not supported.") 
-                    output_values.append(processed_item)  
+                        raise ValueError(f"Embodiment ID {elem['embodiment_id']} not supported.")
+                    output_values.append(processed_item)
                 except (ValueError, SyntaxError, TypeError):
                     # If parsing fails or item is already a string, use it directly
                     if num_views > 1 and elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.AGIBOT.value]:
@@ -150,8 +152,10 @@ def collate(features: List[dict], tokenizer: AutoTokenizer, num_views=3, embodim
                         item = "A multi-view video shows that a robot " + str(item).lower() + " The video is split into four views: The top-left view shows the top camera, the top-right view shows the right camera, the bottom-left view shows the left camera, and the bottom-right view is a black screen. The robot " + str(item).lower()
                     elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.TREX.value]:
                         item = "A multi-view video shows that a bimanual robot with dexterous hands " + str(item).lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right wrist, the bottom-left view shows the camera view from the left wrist, and the bottom-right view is a black screen. The robot " + str(item).lower()
+                    elif elem["embodiment_id"] == embodiment_tag_mapping[EmbodimentTag.ROBOCASA_BIMANUAL_PANDA_INSPIRE_HAND.value]:
+                        item = "A multi-view video shows that a bimanual robot with dexterous hands " + str(item).lower() + " The video is split into four views: The top-left view shows the camera view from the robot's head, the top-right view shows the camera view from the right wrist, the bottom-left view shows the camera view from the left wrist, and the bottom-right view is a black screen. The robot " + str(item).lower()
                     else:
-                        raise ValueError(f"Embodiment ID {elem['embodiment_id']} not supported.")   
+                        raise ValueError(f"Embodiment ID {elem['embodiment_id']} not supported.")
                     output_values.append(item)
             # print("output_values", output_values)
             ids, mask = tokenizer(output_values, return_mask=True, add_special_tokens=True)