feat(recap): add advantage scoring annotation module

Implement the RECAP advantage scoring module as a new phase in
lerobot-annotate. Uses a frozen distributional VF to compute per-frame
advantages, binarizes into positive/negative indicators with per-task
threshold, and writes style=advantage persistent rows for policy
conditioning. Skips VF inference on intervention frames as an optimization.
This commit is contained in:
Khalil Meftah
2026-06-22 14:01:58 +02:00
parent b90ccd283b
commit cc782e3589
11 changed files with 632 additions and 7 deletions
@@ -15,20 +15,24 @@
# limitations under the License.
"""In-process executor that runs the annotation phases.
The executor runs **six phases** in dependency order:
The executor runs **seven phases** in dependency order:
phase 1: ``plan`` module (plan + subtasks + memory)
phase 2: ``interjections`` module (interjections + speech)
phase 3: ``plan`` plan-update pass — re-runs plan emission at every
interjection timestamp produced by phase 2
phase 4: ``vqa`` module (VQA)
phase 5: validator
phase 6: writer
phase 5: ``advantage`` module (advantage scoring via frozen VF)
phase 6: validator
phase 7: writer
Phase 3 is why the ``plan`` module must be re-entered after the
``interjections`` module — to refresh ``plan`` rows at interjection
timestamps.
Phase 5 (advantage) does not depend on the VLM modules, it uses a frozen
distributional value function to compute per-frame advantage indicators.
Distributed execution is provided by Hugging Face Jobs (see
``examples/annotations/run_hf_job.py``); the runner inside the job
invokes ``lerobot-annotate`` which uses this in-process executor.
@@ -74,7 +78,7 @@ class PipelineRunSummary:
@dataclass
class Executor:
"""Run all six phases over a dataset root in-process.
"""Run all seven phases over a dataset root in-process.
Episode-level concurrency comes from ``ExecutorConfig.episode_parallelism``
(a thread pool); cluster-level concurrency comes from running this
@@ -86,6 +90,7 @@ class Executor:
plan: Any # PlanSubtasksMemoryModule
interjections: Any # InterjectionsAndSpeechModule
vqa: Any # GeneralVqaModule
advantage: Any # AdvantageModule
writer: LanguageColumnsWriter
validator: StagingValidator
@@ -112,6 +117,8 @@ class Executor:
phases.append(self._run_plan_update_phase(records, staging_dir))
# Phase 4: ``vqa`` module (VQA)
phases.append(self._run_module_phase("vqa", records, staging_dir, self.vqa))
# Phase 5: ``advantage`` module (advantage scoring via frozen VF)
phases.append(self._run_module_phase("advantage", records, staging_dir, self.advantage))
print("[annotate] running validator...", flush=True)
report = self.validator.validate(records, staging_dir)