From 2686450d68fb042657934006c662ccfdad06b7db Mon Sep 17 00:00:00 2001 From: pepijn Date: Tue, 26 May 2026 05:14:23 +0000 Subject: [PATCH] annotate(plan): force composite-action subtasks; tune run_hf_job for robocasa_smoke MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Subtask prompt (``module_1_subtasks.txt``): - Lock the verb vocabulary to composite atomic actions (``pick up``, ``put``/``place``, ``push``/``pull``, ``turn``, ``press``, ``open``/ ``close``, ``pour``, ``insert``, ``go to``). - Add an explicit ``Forbidden ultra-fine splits`` block instructing the VLM to fold ``move to X`` / ``reach for X`` / ``grasp X`` / ``lift X`` / ``release X`` into the parent composite. Previous examples actively encouraged the over-segmentation pattern. - Rewrite the Good/Bad examples around the composite contract. Job config (``examples/annotations/run_hf_job.py``): - Point at ``pepijn223/robocasa_smoke_2atomic_v3`` on ``h200x4``. - ``--vlm.camera_key=robot0_agentview_left`` (real key for the dataset; the prior ``observation.images.wrist`` did not exist and would have silenced the VQA module). - ``--vlm.serve_command`` ``--max-model-len 131072`` (4x): keeps 90 s @ 1 Hz episode video blocks under context even at full Qwen vision resolution. On 1x H200 (144 GB) the 35B-FP8 model has plenty of room for the bigger KV cache. - ``--vocabulary.enabled=false`` — heterogeneous dataset, no benefit from a single canonical vocabulary. - ``--plan.derive_task_from_video=off``, ``--plan.n_task_rephrasings=0`` — reuse the dataset's own ``episode_task`` strings as-is. - ``--plan.min_subtask_seconds=3.0``, ``--plan.plan_max_steps=6`` — give the new composite-action rules room to land (1.5 s floor was too small to host a full grasp-or-place composite). - ``--vqa.vqa_emission_hz=3.0`` — denser VQA grounding. - Timeout 24h, episode_parallelism=64, client_concurrency=256 to scale to the 25k-trajectory regime when the same recipe is pointed at a larger dataset. Co-authored-by: Cursor --- examples/annotations/run_hf_job.py | 107 +++++++++++------- .../prompts/module_1_subtasks.txt | 58 ++++++++-- 2 files changed, 110 insertions(+), 55 deletions(-) diff --git a/examples/annotations/run_hf_job.py b/examples/annotations/run_hf_job.py index 76d18bb68..be48185f8 100644 --- a/examples/annotations/run_hf_job.py +++ b/examples/annotations/run_hf_job.py @@ -1,15 +1,16 @@ #!/usr/bin/env python """Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6 MoE). -Spawns one ``h200x2`` job that: +Spawns one ``h200x4`` job that: 1. installs this branch of ``lerobot`` plus the annotation extras, - 2. boots two vllm servers (one per GPU) with Qwen3.6-35B-A3B-FP8, - 3. discovers the dataset's canonical subtask + memory vocabulary - from the first 3 sample episodes (phase 0), - 4. runs the plan / interjections / vqa modules across the dataset - (subtasks + memory are constrained to the canonical vocabulary), - 5. uploads the annotated dataset to ``--dest_repo_id`` (when set) + 2. boots four vllm servers (one per H200) with Qwen3.6-35B-A3B-FP8, + 3. runs the plan + vqa modules across the dataset in free-form + mode — phase 0 (canonical vocabulary discovery) is disabled so + every episode's subtasks + memory are generated independently; + interjections is also disabled, which short-circuits the + plan_update phase that depends on it, + 4. uploads the annotated dataset to ``--dest_repo_id`` (when set) or back to ``--repo_id``. Usage: @@ -37,60 +38,80 @@ CMD = ( "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && " "export VLLM_VIDEO_BACKEND=pyav && " "lerobot-annotate " - "--repo_id=imstevenpmwork/super_poulain_draft " - "--dest_repo_id=pepijn223/super_poulain_vocab " + "--repo_id=pepijn223/robocasa_smoke_2atomic_v3 " + "--dest_repo_id=pepijn223/robocasa_smoke_2atomic_v3_annotated " "--push_to_hub=true " "--vlm.backend=openai " "--vlm.model_id=Qwen/Qwen3.6-35B-A3B-FP8 " - "--vlm.parallel_servers=2 " - "--vlm.num_gpus=2 " + "--vlm.parallel_servers=4 " + "--vlm.num_gpus=4 " '--vlm.serve_command="vllm serve Qwen/Qwen3.6-35B-A3B-FP8 ' - "--tensor-parallel-size 1 --max-model-len 32768 " - '--gpu-memory-utilization 0.8 --uvicorn-log-level warning --port {port}" ' + # 4× the context (32768 → 131072) so long episodes at 1 Hz fit even + # at full Qwen vision resolution: 90 frames @ ~700 vision tokens/frame + # ≈ 63 k tokens, comfortably under 131 k. On 1× H200 (144 GB) the + # 35B-FP8 model leaves plenty of room for the bigger KV cache. + "--tensor-parallel-size 1 --max-model-len 131072 " + '--gpu-memory-utilization 0.85 --uvicorn-log-level warning --port {port}" ' "--vlm.serve_ready_timeout_s=1800 " - "--vlm.client_concurrency=128 " + "--vlm.client_concurrency=256 " "--vlm.max_new_tokens=512 " "--vlm.temperature=0.7 " - "--executor.episode_parallelism=16 " + "--executor.episode_parallelism=64 " "--vlm.chat_template_kwargs='{\"enable_thinking\": false}' " - "--vlm.camera_key=observation.images.wrist " -<<<<<<< HEAD:examples/annotation/run_hf_job.py - "--module_1.frames_per_second=1.0 " - "--module_1.use_video_url=true " - "--module_1.use_video_url_fps=1.0 " - "--module_1.derive_task_from_video=always " - "--module_1.n_task_rephrasings=30 " - "--module_2.max_interjections_per_episode=6 " - "--module_3.K=3 " - "--module_3.vqa_emission_hz=1.0 " - "--push_to_hub=pepijn223/super_poulain_full_tool3" -======= - # Phase 0 — canonical vocabulary discovery from the first N sample - # episodes. The VLM picks the right number of subtask + memory - # entries itself from what it sees; the resulting - # meta/canonical_vocabulary.json constrains every subtask + memory - # string to a small repeatable target distribution. - "--vocabulary.sample_episodes=3 " + # Whole-scene agentview is the right choice for subtask reasoning + + # VQA on robocasa: the wrist (``robot0_eye_in_hand``) usually only + # sees the gripper + nearby object, which hurts "what is happening + # in this episode" decomposition. Override per-dataset if your + # cameras are named differently (inspect ``meta/info.json``). + "--vlm.camera_key=observation.images.robot0_agentview_left " + # Phase 0 — canonical vocabulary discovery DISABLED. This dataset's + # episodes span heterogeneous tasks/scenes, so a single shared + # subtask + memory vocabulary would be too narrow — each episode + # generates its subtasks + memory free-form instead. + "--vocabulary.enabled=false " # Phase 1 — plan module (subtasks + plan + memory + task_aug). + "--plan.enabled=true " "--plan.frames_per_second=1.0 " "--plan.use_video_url=true " "--plan.use_video_url_fps=1.0 " - "--plan.derive_task_from_video=always " - "--plan.n_task_rephrasings=30 " - # Phase 2 — interjections + speech. - "--interjections.max_interjections_per_episode=6 " - # Phase 4 — general VQA. - "--vqa.K=3 " - "--vqa.vqa_emission_hz=1.0" ->>>>>>> origin/feat/language-annotation-pipeline:examples/annotations/run_hf_job.py + # Force coarse, composite subtasks (``pick up X`` = approach + grasp + # + lift in one span, not three). 3 s is large enough to host a + # full grasp-or-place composite at typical 20 fps robocasa speeds; + # any candidate span shorter than this gets merged into a neighbour + # by the prompt's authoring rules (see module_1_subtasks.txt). + "--plan.min_subtask_seconds=3.0 " + # Cap so the VLM can't drift into micro-segmentation. Combined with + # the composite-action rules in the prompt, this targets ~3-6 + # meaningful spans per episode for typical pick-and-place demos. + "--plan.plan_max_steps=9 " + # ``off`` keeps the dataset's canonical ``record.episode_task`` as-is + # — no per-episode VLM "what is this video about" call. Switch to + # ``if_short`` (default) only if some episodes have placeholder / + # missing canonical tasks; ``always`` overrides every episode's task. + "--plan.derive_task_from_video=off " + # 0 disables the task_aug pass entirely (see PlanConfig.n_task_rephrasings + # docstring) — no per-episode paraphrase generation, no task_aug rows. + "--plan.n_task_rephrasings=0 " + # Phase 2 — interjections OFF (also skips phase 3 plan_update, + # see executor.py:_run_plan_update_phase guard). + "--interjections.enabled=false " + # Phase 4 — general VQA. K=1 keeps each VQA answer on its own + # emission frame (no temporal smear); see VqaConfig.K docstring. + # 3 Hz cadence: at 20 fps source, that's a VQA tick every ~7 frames. + # NOTE: VQA emits per-camera, so for robocasa (3 cameras) each tick + # produces 3 (user, assistant) row pairs — total call volume ~= 3 * + # 3 Hz * mean_episode_seconds * n_episodes. + "--vqa.enabled=true " + "--vqa.K=1 " + "--vqa.vqa_emission_hz=3.0" ) job = run_job( image="vllm/vllm-openai:latest", command=["bash", "-c", CMD], - flavor="h200x2", + flavor="h200x4", secrets={"HF_TOKEN": token}, - timeout="2h", + timeout="24h", ) print(f"Job URL: {job.url}") print(f"Job ID: {job.id}") diff --git a/src/lerobot/annotations/steerable_pipeline/prompts/module_1_subtasks.txt b/src/lerobot/annotations/steerable_pipeline/prompts/module_1_subtasks.txt index 12bbcfba2..9314282be 100644 --- a/src/lerobot/annotations/steerable_pipeline/prompts/module_1_subtasks.txt +++ b/src/lerobot/annotations/steerable_pipeline/prompts/module_1_subtasks.txt @@ -8,14 +8,42 @@ the robot performs. {vocabulary_block}Authoring rules — Hi Robot atom granularity, pi0.7-style short prompts: -- Each subtask = one atomic skill the low-level policy can execute. -- Write each subtask as an IMPERATIVE COMMAND, starting with a verb: - move, reach, pick up, grasp, place, put, push, pull, open, close, - turn, press, lift, insert, pour... +- Each subtask = one COMPOSITE atomic skill the low-level policy can + execute end-to-end. A "skill" bundles its own approach motion with + its terminal action — do NOT split the approach off as its own + subtask. The whole-arm policy already learns to reach as part of + every manipulation primitive. +- Write each subtask as an IMPERATIVE COMMAND, starting with one of + these verbs (extend only when none fits): + pick up — approach + grasp + lift in one subtask + put on/in — transport + release in one subtask + place on/in — synonym of "put"; pick one and stay consistent + push — contact + linear shove + pull — contact + linear retract + turn — rotary actuation + press