add more changes

examples/dataset/annotate_pgen.py

@@ -1402,6 +1402,13 @@ def main():
         action="store_true",
         help="Push modified dataset to HuggingFace Hub",
     )
+    # add image key
+    parser.add_argument(
+        "--image-key",
+        type=str,
+        default=None,
+        help="Image observation key to use for image mode (default: None)",
+    )
     
     args = parser.parse_args()
     console = Console()
@@ -1443,7 +1450,10 @@ def main():
     )
     
     # Get image keys (for image mode)
-    image_keys = dataset.meta.camera_keys[:args.num_image_views_per_sample]
+    if args.image_key:
+        image_keys = [args.image_key]
+    else:
+        image_keys = dataset.meta.camera_keys[:args.num_image_views_per_sample]
     if not args.video_mode:
         console.print(f"[cyan]Using image keys: {image_keys}[/cyan]")
     

examples/dataset/run.sh

@@ -1,10 +1,11 @@
+python examples/dataset/annotate.py \
+    --repo-id jadechoghari/collect-data \
+    --video-key observation.images.base \
+    --model Qwen/Qwen3-VL-30B-A3B-Instruct \
+    --episodes 16 22
+
 # python examples/dataset/annotate.py \
 #     --repo-id lerobot/svla_so101_pickplace \
 #     --video-key observation.images.side \
 #     --model Qwen/Qwen3-VL-30B-A3B-Instruct \
-
-python examples/dataset/annotate.py \
-    --repo-id lerobot/svla_so101_pickplace \
-    --video-key observation.images.side \
-    --model Qwen/Qwen3-VL-30B-A3B-Instruct \
-    --episodes 5
+#     --episodes 5

examples/dataset/run_pgen.sh

@@ -4,12 +4,12 @@
 # This generates user prompts and robot utterances for hierarchical policy training
 
 # Configuration
-REPO_ID="lerobot/svla_so101_pickplace"
+REPO_ID="jadechoghari/collect-data"
 MODEL="Qwen/Qwen3-VL-30B-A3B-Instruct"
 # Alternative: MODEL="Qwen/Qwen2-VL-7B-Instruct"
 
 
-OUTPUT_DIR="/fsx/jade_choghari/outputs/pgen_annotations1"
+OUTPUT_DIR="/fsx/jade_choghari/outputs/collect-data-pgen"
 BATCH_SIZE=32
 TEMPERATURE=0.9
 SAMPLE_INTERVAL=5.0  # Generate dialogue every 1 second (all episodes processed)
@@ -22,6 +22,7 @@ python examples/dataset/annotate_pgen.py \
     --temperature "$TEMPERATURE" \
     --batch-size "$BATCH_SIZE" \
     --sample-interval "$SAMPLE_INTERVAL" \
+    --image-key observation.images.base \
     --num-image-views-per-sample 1
 
 # For faster testing, increase sample interval: