test: add dataset guard + fix imports

Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>

.github/workflows/benchmark_tests.yml

@@ -83,13 +83,10 @@ jobs:
           cache-binary: false
 
       - name: Login to Docker Hub
-        if: ${{ env.DOCKERHUB_USERNAME != '' }}
         uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
         with:
           username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
           password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
-        env:
-          DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
 
       # Build the benchmark-specific image. The Dockerfile separates dep-install
       # from source-copy, so code-only changes skip the slow uv-sync layer
@@ -118,7 +115,7 @@ jobs:
             bash -c "
               hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
               lerobot-eval \
-                --policy.path=lerobot/smolvla_libero \
+                --policy.path=pepijn223/smolvla_libero \
                 --env.type=libero \
                 --env.task=libero_spatial \
                 --eval.batch_size=1 \
@@ -147,7 +144,7 @@ jobs:
             --artifacts-dir /tmp/libero-artifacts \
             --env libero \
             --task libero_spatial \
-            --policy lerobot/smolvla_libero
+            --policy pepijn223/smolvla_libero
 
       - name: Upload Libero rollout video
         if: always()
@@ -241,13 +238,10 @@ jobs:
           cache-binary: false
 
       - name: Login to Docker Hub
-        if: ${{ env.DOCKERHUB_USERNAME != '' }}
         uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
         with:
           username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
           password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
-        env:
-          DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
 
       - name: Build MetaWorld benchmark image
         uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
@@ -270,7 +264,7 @@ jobs:
             bash -c "
               hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
               lerobot-eval \
-                --policy.path=lerobot/smolvla_metaworld \
+                --policy.path=pepijn223/smolvla_metaworld \
                 --env.type=metaworld \
                 --env.task=metaworld-push-v3 \
                 --eval.batch_size=1 \
@@ -299,7 +293,7 @@ jobs:
             --artifacts-dir /tmp/metaworld-artifacts \
             --env metaworld \
             --task metaworld-push-v3 \
-            --policy lerobot/smolvla_metaworld
+            --policy pepijn223/smolvla_metaworld
 
       - name: Upload MetaWorld rollout video
         if: always()
@@ -316,630 +310,3 @@ jobs:
           name: metaworld-metrics
           path: /tmp/metaworld-artifacts/metrics.json
           if-no-files-found: warn
-
-  # ── ROBOTWIN 2.0 ──────────────────────────────────────────────────────────
-  # Isolated image: full RoboTwin 2.0 stack — SAPIEN, mplib, CuRobo,
-  # pytorch3d, + simulation assets (~4 GB).
-  # Build takes ~20 min on first run; subsequent runs hit the layer cache.
-  # Requires an NVIDIA GPU runner with CUDA 12.1 drivers.
-  robotwin-integration-test:
-    name: RoboTwin 2.0 — build image + 1-episode eval
-    runs-on:
-      group: aws-g6-4xlarge-plus
-    env:
-      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
-      ROBOTWIN_POLICY: lerobot/smolvla_robotwin
-      ROBOTWIN_TASKS: beat_block_hammer,click_bell,handover_block,stack_blocks_two,click_alarmclock,open_microwave,adjust_bottle,lift_pot,stamp_seal,turn_switch
-
-    steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
-        with:
-          persist-credentials: false
-          lfs: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          cache-binary: false
-
-      - name: Login to Docker Hub
-        if: ${{ env.DOCKERHUB_USERNAME != '' }}
-        uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
-        env:
-          DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-
-      # Build the full-install image: SAPIEN, mplib, CuRobo, pytorch3d +
-      # simulation assets (~4 GB). Layer cache lives in the runner's local
-      # Docker daemon — reused across re-runs on the same machine.
-      - name: Build RoboTwin 2.0 benchmark image
-        uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
-        with:
-          context: .
-          file: docker/Dockerfile.benchmark.robotwin
-          push: false
-          load: true
-          tags: lerobot-benchmark-robotwin:ci
-          cache-from: type=local,src=/tmp/.buildx-cache-robotwin
-          cache-to: type=local,dest=/tmp/.buildx-cache-robotwin,mode=max
-
-      - name: Run RoboTwin 2.0 smoke eval (10 tasks, 1 episode each)
-        if: env.HF_USER_TOKEN != ''
-        run: |
-          # Named container (no --rm) so we can docker cp artifacts out.
-          docker run --name robotwin-eval --gpus all \
-            --shm-size=4g \
-            -e HF_HOME=/tmp/hf \
-            -e HF_USER_TOKEN="${HF_USER_TOKEN}" \
-            -e ROBOTWIN_POLICY="${ROBOTWIN_POLICY}" \
-            -e ROBOTWIN_TASKS="${ROBOTWIN_TASKS}" \
-            lerobot-benchmark-robotwin:ci \
-            bash -c "
-              hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
-              cd /opt/robotwin && lerobot-eval \
-                --policy.path=\"\$ROBOTWIN_POLICY\" \
-                --env.type=robotwin \
-                --env.task=\"\$ROBOTWIN_TASKS\" \
-                --eval.batch_size=1 \
-                --eval.n_episodes=1 \
-                --eval.use_async_envs=false \
-                --policy.device=cuda \
-                '--rename_map={\"observation.images.head_camera\": \"observation.images.camera1\", \"observation.images.left_camera\": \"observation.images.camera2\", \"observation.images.right_camera\": \"observation.images.camera3\"}' \
-                --output_dir=/tmp/eval-artifacts
-              python /lerobot/scripts/ci/extract_task_descriptions.py \
-                --env robotwin \
-                --task \"\$ROBOTWIN_TASKS\" \
-                --output /tmp/eval-artifacts/task_descriptions.json
-            "
-
-      - name: Copy RoboTwin artifacts from container
-        if: always()
-        run: |
-          mkdir -p /tmp/robotwin-artifacts
-          docker cp robotwin-eval:/tmp/eval-artifacts/. /tmp/robotwin-artifacts/ 2>/dev/null || true
-          docker rm -f robotwin-eval || true
-
-      - name: Parse RoboTwin eval metrics
-        if: always()
-        run: |
-          python3 scripts/ci/parse_eval_metrics.py \
-            --artifacts-dir /tmp/robotwin-artifacts \
-            --env robotwin \
-            --task "${ROBOTWIN_TASKS}" \
-            --policy "${ROBOTWIN_POLICY}"
-
-      - name: Upload RoboTwin rollout video
-        if: always()
-        uses: actions/upload-artifact@v4
-        with:
-          name: robotwin-rollout-video
-          path: /tmp/robotwin-artifacts/videos/
-          if-no-files-found: warn
-
-      - name: Upload RoboTwin eval metrics
-        if: always()
-        uses: actions/upload-artifact@v4
-        with:
-          name: robotwin-metrics
-          path: /tmp/robotwin-artifacts/metrics.json
-          if-no-files-found: warn
-
-  # ── ROBOCASA365 ──────────────────────────────────────────────────────────
-  # Isolated image: robocasa + robosuite installed manually as editable
-  # clones (no `lerobot[robocasa]` extra — robocasa's setup.py pins
-  # `lerobot==0.3.3`, which would shadow this repo's lerobot).
-  robocasa-integration-test:
-    name: RoboCasa365 — build image + 1-episode eval
-    runs-on:
-      group: aws-g6-4xlarge-plus
-    env:
-      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
-
-    steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
-        with:
-          persist-credentials: false
-          lfs: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          cache-binary: false
-
-      - name: Login to Docker Hub
-        if: ${{ env.DOCKERHUB_USERNAME != '' }}
-        uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
-        env:
-          DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-
-      - name: Build RoboCasa365 benchmark image
-        uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
-        with:
-          context: .
-          file: docker/Dockerfile.benchmark.robocasa
-          push: false
-          load: true
-          tags: lerobot-benchmark-robocasa:ci
-
-      - name: Run RoboCasa365 smoke eval (10 atomic tasks, 1 episode each)
-        if: env.HF_USER_TOKEN != ''
-        run: |
-          docker run --name robocasa-eval --gpus all \
-            --shm-size=4g \
-            -e HF_HOME=/tmp/hf \
-            -e HF_USER_TOKEN="${HF_USER_TOKEN}" \
-            -e HF_HUB_DOWNLOAD_TIMEOUT=300 \
-            -e MUJOCO_GL=egl \
-            lerobot-benchmark-robocasa:ci \
-            bash -c "
-              hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
-              lerobot-eval \
-                --policy.path=lerobot/smolvla_robocasa \
-                --env.type=robocasa \
-                --env.task=CloseFridge,OpenCabinet,OpenDrawer,TurnOnMicrowave,TurnOffStove,CloseToasterOvenDoor,SlideDishwasherRack,TurnOnSinkFaucet,NavigateKitchen,TurnOnElectricKettle \
-                --eval.batch_size=1 \
-                --eval.n_episodes=1 \
-                --eval.use_async_envs=false \
-                --policy.device=cuda \
-                '--rename_map={\"observation.images.robot0_agentview_left\": \"observation.images.camera1\", \"observation.images.robot0_eye_in_hand\": \"observation.images.camera2\", \"observation.images.robot0_agentview_right\": \"observation.images.camera3\"}' \
-                --output_dir=/tmp/eval-artifacts
-              python scripts/ci/extract_task_descriptions.py \
-                --env robocasa \
-                --task CloseFridge,OpenCabinet,OpenDrawer,TurnOnMicrowave,TurnOffStove,CloseToasterOvenDoor,SlideDishwasherRack,TurnOnSinkFaucet,NavigateKitchen,TurnOnElectricKettle \
-                --output /tmp/eval-artifacts/task_descriptions.json
-            "
-
-      - name: Copy RoboCasa365 artifacts from container
-        if: always()
-        run: |
-          mkdir -p /tmp/robocasa-artifacts
-          docker cp robocasa-eval:/tmp/eval-artifacts/. /tmp/robocasa-artifacts/ 2>/dev/null || true
-          docker rm -f robocasa-eval || true
-
-      - name: Parse RoboCasa365 eval metrics
-        if: always()
-        run: |
-          python3 scripts/ci/parse_eval_metrics.py \
-            --artifacts-dir /tmp/robocasa-artifacts \
-            --env robocasa \
-            --task atomic_smoke_10 \
-            --policy lerobot/smolvla_robocasa
-
-      - name: Upload RoboCasa365 rollout video
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: robocasa-rollout-video
-          path: /tmp/robocasa-artifacts/videos/
-          if-no-files-found: warn
-
-      - name: Upload RoboCasa365 eval metrics
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: robocasa-metrics
-          path: /tmp/robocasa-artifacts/metrics.json
-          if-no-files-found: warn
-
-  # ── ROBOCEREBRA ───────────────────────────────────────────────────────────
-  # Reuses the LIBERO simulator (libero_10 suite) with RoboCerebra camera
-  # defaults (image/wrist_image). The image is layered on
-  # huggingface/lerobot-gpu, which already ships [libero] as part of [all].
-  robocerebra-integration-test:
-    name: RoboCerebra — build image + 1-episode eval
-    runs-on:
-      group: aws-g6-4xlarge-plus
-    env:
-      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
-
-    steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
-        with:
-          persist-credentials: false
-          lfs: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          cache-binary: false
-
-      - name: Login to Docker Hub
-        if: ${{ env.DOCKERHUB_USERNAME != '' }}
-        uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
-        env:
-          DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-
-      - name: Build RoboCerebra benchmark image
-        uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
-        with:
-          context: .
-          file: docker/Dockerfile.benchmark.robocerebra
-          push: false
-          load: true
-          tags: lerobot-benchmark-robocerebra:ci
-          cache-from: type=local,src=/tmp/.buildx-cache-robocerebra
-          cache-to: type=local,dest=/tmp/.buildx-cache-robocerebra,mode=max
-
-      - name: Run RoboCerebra smoke eval (1 episode)
-        if: env.HF_USER_TOKEN != ''
-        run: |
-          docker run --name robocerebra-eval --gpus all \
-            --shm-size=4g \
-            -e HF_HOME=/tmp/hf \
-            -e HF_USER_TOKEN="${HF_USER_TOKEN}" \
-            -e HF_HUB_DOWNLOAD_TIMEOUT=300 \
-            -e LIBERO_DATA_FOLDER=/tmp/libero_data \
-            lerobot-benchmark-robocerebra:ci \
-            bash -c "
-              hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
-              lerobot-eval \
-                --policy.path=lerobot/smolvla_robocerebra \
-                --env.type=libero \
-                --env.task=libero_10 \
-                --env.fps=20 \
-                --env.obs_type=pixels_agent_pos \
-                --env.observation_height=256 \
-                --env.observation_width=256 \
-                '--env.camera_name_mapping={\"agentview_image\": \"image\", \"robot0_eye_in_hand_image\": \"wrist_image\"}' \
-                --eval.batch_size=1 \
-                --eval.n_episodes=1 \
-                --eval.use_async_envs=false \
-                --policy.device=cuda \
-                '--rename_map={\"observation.images.image\": \"observation.images.camera1\", \"observation.images.wrist_image\": \"observation.images.camera2\"}' \
-                --policy.empty_cameras=1 \
-                --output_dir=/tmp/eval-artifacts
-              python scripts/ci/extract_task_descriptions.py \
-                --env libero --task libero_10 \
-                --output /tmp/eval-artifacts/task_descriptions.json
-            "
-
-      - name: Copy RoboCerebra artifacts from container
-        if: always()
-        run: |
-          mkdir -p /tmp/robocerebra-artifacts
-          docker cp robocerebra-eval:/tmp/eval-artifacts/. /tmp/robocerebra-artifacts/ 2>/dev/null || true
-          docker rm -f robocerebra-eval || true
-
-      - name: Parse RoboCerebra eval metrics
-        if: always()
-        run: |
-          python3 scripts/ci/parse_eval_metrics.py \
-            --artifacts-dir /tmp/robocerebra-artifacts \
-            --env robocerebra \
-            --task libero_10 \
-            --policy lerobot/smolvla_robocerebra
-
-      - name: Upload RoboCerebra rollout video
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: robocerebra-rollout-video
-          path: /tmp/robocerebra-artifacts/videos/
-          if-no-files-found: warn
-
-      - name: Upload RoboCerebra eval metrics
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: robocerebra-metrics
-          path: /tmp/robocerebra-artifacts/metrics.json
-          if-no-files-found: warn
-
-  # ── ROBOMME ───────────────────────────────────────────────────────────────
-  # Isolated image: mani-skill/SAPIEN/Vulkan chain with gymnasium and numpy
-  # overrides (robomme can't be a pyproject extra due to numpy<2 pin).
-  robomme-integration-test:
-    name: RoboMME — build image + 1-episode eval
-    runs-on:
-      group: aws-g6-4xlarge-plus
-    env:
-      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
-      ROBOMME_POLICY: lerobot/smolvla_robomme
-      ROBOMME_TASKS: PickXtimes,BinFill,StopCube,MoveCube,InsertPeg,SwingXtimes,VideoUnmask,ButtonUnmask,PickHighlight,PatternLock
-
-    steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
-        with:
-          persist-credentials: false
-          lfs: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          cache-binary: false
-
-      - name: Login to Docker Hub
-        if: ${{ env.DOCKERHUB_USERNAME != '' }}
-        uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
-        env:
-          DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-
-      - name: Build RoboMME benchmark image
-        uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
-        with:
-          context: .
-          file: docker/Dockerfile.benchmark.robomme
-          push: false
-          load: true
-          tags: lerobot-benchmark-robomme:ci
-
-      - name: Run RoboMME smoke eval (10 tasks, 1 episode each)
-        if: env.HF_USER_TOKEN != ''
-        run: |
-          docker run --name robomme-eval --gpus all \
-            --shm-size=4g \
-            -e HF_HOME=/tmp/hf \
-            -e HF_USER_TOKEN="${HF_USER_TOKEN}" \
-            -e HF_HUB_DOWNLOAD_TIMEOUT=300 \
-            -e ROBOMME_POLICY="${ROBOMME_POLICY}" \
-            -e ROBOMME_TASKS="${ROBOMME_TASKS}" \
-            lerobot-benchmark-robomme:ci \
-            bash -c "
-              hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
-              lerobot-eval \
-                --policy.path=\"\$ROBOMME_POLICY\" \
-                --env.type=robomme \
-                --env.task=\"\$ROBOMME_TASKS\" \
-                --env.dataset_split=test \
-                --env.task_ids=[0] \
-                --eval.batch_size=1 \
-                --eval.n_episodes=1 \
-                --eval.use_async_envs=false \
-                --policy.device=cuda \
-                '--rename_map={\"observation.images.image\": \"observation.images.camera1\", \"observation.images.wrist_image\": \"observation.images.camera2\"}' \
-                --policy.empty_cameras=3 \
-                --output_dir=/tmp/eval-artifacts
-              python scripts/ci/extract_task_descriptions.py \
-                --env robomme --task \"\$ROBOMME_TASKS\" \
-                --output /tmp/eval-artifacts/task_descriptions.json
-            "
-
-      - name: Copy RoboMME artifacts from container
-        if: always()
-        run: |
-          mkdir -p /tmp/robomme-artifacts
-          docker cp robomme-eval:/tmp/eval-artifacts/. /tmp/robomme-artifacts/ 2>/dev/null || true
-          docker rm -f robomme-eval || true
-
-      - name: Parse RoboMME eval metrics
-        if: always()
-        run: |
-          python3 scripts/ci/parse_eval_metrics.py \
-            --artifacts-dir /tmp/robomme-artifacts \
-            --env robomme \
-            --task "${ROBOMME_TASKS}" \
-            --policy "${ROBOMME_POLICY}"
-
-      - name: Upload RoboMME rollout video
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: robomme-rollout-video
-          path: /tmp/robomme-artifacts/videos/
-          if-no-files-found: warn
-
-      - name: Upload RoboMME eval metrics
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: robomme-metrics
-          path: /tmp/robomme-artifacts/metrics.json
-          if-no-files-found: warn
-
-  # ── LIBERO-plus ───────────────────────────────────────────────────────────
-  # Isolated image: LIBERO-plus fork cloned into /home/user_lerobot on top of
-  # huggingface/lerobot-gpu (see docker/Dockerfile.benchmark.libero_plus).
-  libero-plus-integration-test:
-    name: LIBERO-plus — build image + 1-episode eval
-    runs-on:
-      group: aws-g6-4xlarge-plus
-    env:
-      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
-      LIBERO_PLUS_SUITE: libero_spatial
-      LIBERO_PLUS_POLICY: lerobot/smolvla_libero_plus
-      LIBERO_PLUS_TASK_IDS: "[0,100,260,500,1000,1500,2000,2400]"
-
-    steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
-        with:
-          persist-credentials: false
-          lfs: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          cache-binary: false
-
-      - name: Login to Docker Hub
-        if: ${{ env.DOCKERHUB_USERNAME != '' }}
-        uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
-        env:
-          DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-
-      - name: Build LIBERO-plus benchmark image
-        uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
-        with:
-          context: .
-          file: docker/Dockerfile.benchmark.libero_plus
-          push: false
-          load: true
-          tags: lerobot-benchmark-libero-plus:ci
-          cache-from: type=local,src=/tmp/.buildx-cache-libero-plus
-          cache-to: type=local,dest=/tmp/.buildx-cache-libero-plus,mode=max
-
-      - name: Run LIBERO-plus smoke eval (1 episode)
-        if: env.HF_USER_TOKEN != ''
-        run: |
-          docker run --name libero-plus-eval --gpus all \
-            --shm-size=4g \
-            -e HF_HOME=/tmp/hf \
-            -e HF_USER_TOKEN="${HF_USER_TOKEN}" \
-            -e HF_HUB_DOWNLOAD_TIMEOUT=300 \
-            -e LIBERO_PLUS_SUITE="${LIBERO_PLUS_SUITE}" \
-            -e LIBERO_PLUS_POLICY="${LIBERO_PLUS_POLICY}" \
-            -e LIBERO_PLUS_TASK_IDS="${LIBERO_PLUS_TASK_IDS}" \
-            lerobot-benchmark-libero-plus:ci \
-            bash -c "
-              hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
-              lerobot-eval \
-                --policy.path=\"\$LIBERO_PLUS_POLICY\" \
-                --env.type=libero_plus \
-                --env.task=\"\$LIBERO_PLUS_SUITE\" \
-                --env.task_ids=\"\$LIBERO_PLUS_TASK_IDS\" \
-                --eval.batch_size=1 \
-                --eval.n_episodes=1 \
-                --eval.use_async_envs=false \
-                --policy.device=cuda \
-                '--env.camera_name_mapping={\"agentview_image\": \"camera1\", \"robot0_eye_in_hand_image\": \"camera2\"}' \
-                --policy.empty_cameras=1 \
-                --output_dir=/tmp/eval-artifacts
-              python scripts/ci/extract_task_descriptions.py \
-                --env libero_plus --task \"\$LIBERO_PLUS_SUITE\" \
-                --output /tmp/eval-artifacts/task_descriptions.json
-            "
-
-      - name: Copy LIBERO-plus artifacts from container
-        if: always()
-        run: |
-          mkdir -p /tmp/libero-plus-artifacts
-          docker cp libero-plus-eval:/tmp/eval-artifacts/. /tmp/libero-plus-artifacts/ 2>/dev/null || true
-          docker rm -f libero-plus-eval || true
-
-      - name: Parse LIBERO-plus eval metrics
-        if: always()
-        run: |
-          python3 scripts/ci/parse_eval_metrics.py \
-            --artifacts-dir /tmp/libero-plus-artifacts \
-            --env libero_plus \
-            --task "${LIBERO_PLUS_SUITE}" \
-            --policy "${LIBERO_PLUS_POLICY}"
-
-      - name: Upload LIBERO-plus rollout video
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: libero-plus-rollout-video
-          path: /tmp/libero-plus-artifacts/videos/
-          if-no-files-found: warn
-
-      - name: Upload LIBERO-plus eval metrics
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: libero-plus-metrics
-          path: /tmp/libero-plus-artifacts/metrics.json
-          if-no-files-found: warn
-
-  # ── VLABENCH ─────────────────────────────────────────────────────────────
-  # Isolated image: lerobot[vlabench] only (VLABench, mujoco==3.2.2, dm-control chain)
-  vlabench-integration-test:
-    name: VLABench — build image + 1-episode eval
-    runs-on:
-      group: aws-g6-4xlarge-plus
-    env:
-      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
-
-    steps:
-      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
-        with:
-          persist-credentials: false
-          lfs: true
-
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          cache-binary: false
-
-      - name: Login to Docker Hub
-        if: ${{ env.DOCKERHUB_USERNAME != '' }}
-        uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses]
-        with:
-          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
-        env:
-          DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
-
-      - name: Build VLABench benchmark image
-        uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
-        with:
-          context: .
-          file: docker/Dockerfile.benchmark.vlabench
-          push: false
-          load: true
-          tags: lerobot-benchmark-vlabench:ci
-          build-args: |
-            VLABENCH_ASSETS_REPO=lerobot/vlabench-assets
-
-      - name: Run VLABench smoke eval (10 tasks, 1 episode each)
-        if: env.HF_USER_TOKEN != ''
-        run: |
-          docker run --name vlabench-eval --gpus all \
-            --shm-size=4g \
-            -e HF_HOME=/tmp/hf \
-            -e HF_USER_TOKEN="${HF_USER_TOKEN}" \
-            -e HF_HUB_DOWNLOAD_TIMEOUT=300 \
-            -e MUJOCO_GL=egl \
-            lerobot-benchmark-vlabench:ci \
-            bash -c "
-              hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
-              lerobot-eval \
-                --policy.path=lerobot/smolvla_vlabench \
-                --env.type=vlabench \
-                --env.task=select_fruit,select_toy,select_book,select_painting,select_drink,select_ingredient,select_billiards,select_poker,add_condiment,insert_flower \
-                --eval.batch_size=1 \
-                --eval.n_episodes=1 \
-                --eval.use_async_envs=false \
-                --policy.device=cuda \
-                '--rename_map={\"observation.images.image\": \"observation.images.camera1\", \"observation.images.second_image\": \"observation.images.camera2\", \"observation.images.wrist_image\": \"observation.images.camera3\"}' \
-                --output_dir=/tmp/eval-artifacts
-              python scripts/ci/extract_task_descriptions.py \
-                --env vlabench \
-                --task select_fruit,select_toy,select_book,select_painting,select_drink,select_ingredient,select_billiards,select_poker,add_condiment,insert_flower \
-                --output /tmp/eval-artifacts/task_descriptions.json
-            "
-
-      - name: Copy VLABench artifacts from container
-        if: always()
-        run: |
-          mkdir -p /tmp/vlabench-artifacts
-          docker cp vlabench-eval:/tmp/eval-artifacts/. /tmp/vlabench-artifacts/ 2>/dev/null || true
-          docker rm -f vlabench-eval || true
-
-      - name: Parse VLABench eval metrics
-        if: always()
-        run: |
-          python3 scripts/ci/parse_eval_metrics.py \
-            --artifacts-dir /tmp/vlabench-artifacts \
-            --env vlabench \
-            --task select_fruit,select_toy,select_book,select_painting,select_drink,select_ingredient,select_billiards,select_poker,add_condiment,insert_flower \
-            --policy lerobot/smolvla_vlabench
-
-      - name: Upload VLABench rollout video
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: vlabench-rollout-video
-          path: /tmp/vlabench-artifacts/videos/
-          if-no-files-found: warn
-
-      - name: Upload VLABench eval metrics
-        if: always()
-        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
-        with:
-          name: vlabench-metrics
-          path: /tmp/vlabench-artifacts/metrics.json
-          if-no-files-found: warn

Makefile

@@ -178,9 +178,3 @@ test-smolvla-ete-eval:
 		--env.episode_length=5 \
 		--eval.n_episodes=1 \
 		--eval.batch_size=1
-
-# E2E annotation pipeline smoke test against a tiny in-memory fixture
-# dataset. Opt-in (not part of `make test-end-to-end`) and uses a stub VLM
-# backend, so it does not require a real model checkpoint or GPU.
-annotation-e2e:
-	uv run python -m tests.annotations.run_e2e_smoke

docker/Dockerfile.benchmark.libero_plus

@@ -1,84 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Benchmark image for LIBERO-plus integration tests.
-# Extends the nightly GPU image (which has lerobot[all]) with the LIBERO-plus
-# fork source + its 6.4 GB perturbation assets.
-#
-# Build:  docker build -f docker/Dockerfile.benchmark.libero_plus -t lerobot-benchmark-libero-plus .
-# Run:    docker run --gpus all --rm lerobot-benchmark-libero-plus lerobot-eval ...
-
-FROM huggingface/lerobot-gpu:latest
-ENV MUJOCO_GL=egl
-
-# unzip for the 6.4 GB assets.zip; the rest are LIBERO-plus build-time extras
-# (wand / ImageMagick / fontconfig) not in the nightly base.
-USER root
-RUN apt-get update \
-    && apt-get install -y --no-install-recommends \
-         unzip libexpat1 libfontconfig1-dev libmagickwand-dev \
-    && apt-get clean && rm -rf /var/lib/apt/lists/*
-USER user_lerobot
-
-# robosuite==1.4.1 is mandatory (the fork uses `single_arm_env` removed in
-# v1.5+). The rest are LIBERO-plus runtime deps pulled from its setup.py.
-# We install these explicitly instead of via the [libero_plus] extra because
-# the extra's `libero @ git+...` dep installs as a namespace package and then
-# clone and PYTHONPATH-override it below.
-RUN uv pip install --no-cache \
-        "robosuite==1.4.1" \
-        "bddl==1.0.1" \
-        "easydict==1.13" \
-        "mujoco==3.7.0" \
-        "matplotlib==3.10.8" \
-        "Wand==0.6.13" \
-        "scikit-image==0.25.2" \
-        "gym==0.26.2"
-
-# Clone LIBERO-plus and make it importable as `libero`. The nightly base has
-# hf-libero (10 tasks) preinstalled via lerobot[libero]; uninstall it so
-# Python resolves `import libero` to the 2402-task LIBERO-plus module instead.
-# Pinned to the current upstream main SHA so benchmark builds stay reproducible.
-ARG LIBERO_PLUS_SHA=4976dc3
-ENV LIBERO_PLUS_ROOT=/home/user_lerobot/libero-plus/libero/libero
-RUN git clone https://github.com/sylvestf/LIBERO-plus.git /home/user_lerobot/libero-plus \
-    && git -C /home/user_lerobot/libero-plus checkout ${LIBERO_PLUS_SHA} \
-    && cd /home/user_lerobot/libero-plus && uv pip install --no-cache --no-deps -e "." \
-    && (uv pip uninstall hf-libero 2>/dev/null || true)
-ENV PYTHONPATH="/home/user_lerobot/libero-plus:${PYTHONPATH}"
-
-# Perturbation textures/scenes: bddl_base_domain.py resolves XMLs via
-# DIR_PATH/../assets (package-relative, ignoring ~/.libero/config.yaml). All
-# 2402 tasks reference files that ship only in Sylvest/LIBERO-plus's
-# assets.zip (6.4 GB) under a deep author-internal prefix — extract and
-# flatten it under ${LIBERO_PLUS_ROOT}/assets.
-RUN python -c "\
-from huggingface_hub import hf_hub_download; \
-hf_hub_download(repo_id='Sylvest/LIBERO-plus', repo_type='dataset', \
-                filename='assets.zip', local_dir='/tmp/libero-plus-dl')" \
-    && unzip -q /tmp/libero-plus-dl/assets.zip -d /tmp/libero-plus-dl/extract \
-    && ASSETS_DIR=$(find /tmp/libero-plus-dl/extract -type d -name assets | head -1) \
-    && mv "${ASSETS_DIR}" ${LIBERO_PLUS_ROOT}/assets \
-    && rm -rf /tmp/libero-plus-dl
-
-# Point ~/.libero/config.yaml at the clone so LIBERO-plus's imports are
-# non-interactive (it calls input() when the config is missing).
-RUN mkdir -p /home/user_lerobot/.libero \
-    && printf "assets: ${LIBERO_PLUS_ROOT}/assets\nbddl_files: ${LIBERO_PLUS_ROOT}/bddl_files\ndatasets: ${LIBERO_PLUS_ROOT}/../datasets\ninit_states: ${LIBERO_PLUS_ROOT}/init_files\n" \
-       > /home/user_lerobot/.libero/config.yaml
-
-# Overlay the PR's source code on top of the nightly image.
-COPY --chown=user_lerobot:user_lerobot . .
-
-CMD ["/bin/bash"]

docker/Dockerfile.benchmark.robocasa

@@ -1,71 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Benchmark image for RoboCasa365 integration tests.
-# Extends the nightly GPU image (which already has all extras installed)
-# with the PR's source code and RoboCasa-specific asset setup.
-#
-# Build:  docker build -f docker/Dockerfile.benchmark.robocasa -t lerobot-benchmark-robocasa .
-# Run:    docker run --gpus all --rm lerobot-benchmark-robocasa lerobot-eval ...
-
-FROM huggingface/lerobot-gpu:latest
-
-# Install robocasa + robosuite as editable clones. pip-installing from git
-# omits data files like robocasa/models/assets/box_links/box_links_assets.json
-# (not declared in package_data), which download_kitchen_assets needs at import.
-#
-# `--no-deps` on robocasa is deliberate: its setup.py pins `lerobot==0.3.3`
-# in install_requires, which would shadow the editable lerobot baked into
-# this image. We install robocasa's actual runtime deps explicitly instead.
-# Pinned SHAs for reproducible benchmark runs. Bump when you need an
-# upstream fix; don't rely on `main`/`master` drift.
-ARG ROBOCASA_SHA=56e355ccc64389dfc1b8a61a33b9127b975ba681
-ARG ROBOSUITE_SHA=aaa8b9b214ce8e77e82926d677b4d61d55e577ab
-RUN git clone https://github.com/robocasa/robocasa.git ~/robocasa && \
-    git -C ~/robocasa checkout ${ROBOCASA_SHA} && \
-    git clone https://github.com/ARISE-Initiative/robosuite.git ~/robosuite && \
-    git -C ~/robosuite checkout ${ROBOSUITE_SHA} && \
-    uv pip install --no-cache -e ~/robocasa --no-deps && \
-    uv pip install --no-cache -e ~/robosuite && \
-    uv pip install --no-cache \
-      "numpy==2.2.5" "numba==0.61.2" "scipy==1.15.3" "mujoco==3.3.1" \
-      "pygame==2.6.1" "Pillow==12.2.0" "opencv-python==4.13.0.92" \
-      "pyyaml==6.0.3" "pynput==1.8.1" "tqdm==4.67.3" "termcolor==3.3.0" \
-      "imageio==2.37.3" "h5py==3.16.0" "lxml==6.0.4" "hidapi==0.14.0.post4" \
-      "tianshou==0.4.10" "gymnasium==1.2.3"
-
-# Set up robocasa macros and download kitchen assets. We need:
-#   - tex              : base environment textures
-#   - tex_generative   : AI-generated textures; kitchen fixture XMLs embed
-#                        refs to generative_textures/wall/tex*.png
-#                        unconditionally, so MjModel.from_xml_string fails
-#                        at reset time without them (even if the env is
-#                        constructed with generative_textures=None).
-#   - fixtures_lw      : lightwheel kitchen fixtures (fridge, counters...)
-#   - objs_lw          : lightwheel object meshes (stools, misc props)
-# We skip the objaverse/aigen object packs (~30GB combined) by pairing
-# this with --env.obj_registries=["lightwheel"] on the lerobot side.
-# The download script prompts interactively, so pipe 'y' to auto-accept.
-RUN python -m robocasa.scripts.setup_macros && \
-    yes y | python -m robocasa.scripts.download_kitchen_assets \
-      --type tex tex_generative fixtures_lw objs_lw
-
-# Overlay the PR's source code on top of the nightly image.
-COPY --chown=user_lerobot:user_lerobot . .
-
-# Re-install lerobot editably so the new source (with RoboCasaEnv registration)
-# replaces the stale package baked into the nightly image.
-RUN uv pip install --no-cache --no-deps -e .
-
-CMD ["/bin/bash"]

docker/Dockerfile.benchmark.robocerebra

@@ -1,43 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Benchmark image for RoboCerebra integration tests.
-# RoboCerebra reuses LIBERO's simulator (libero_10 suite) with a different
-# rename_map, so this image is identical to the LIBERO benchmark image —
-# extends the nightly GPU base with LIBERO assets + the PR's source code.
-#
-# Build:  docker build -f docker/Dockerfile.benchmark.robocerebra -t lerobot-benchmark-robocerebra .
-# Run:    docker run --gpus all --rm lerobot-benchmark-robocerebra lerobot-eval ...
-
-FROM huggingface/lerobot-gpu:latest
-
-# Pre-download lerobot/libero-assets from HF Hub so nothing is fetched at
-# runtime (which times out on CI). Point the libero config at the cached path.
-# libero/libero/__init__.py calls input() when ~/.libero/config.yaml is missing,
-# so we write the config before any libero import can happen.
-RUN LIBERO_DIR=$(python -c \
-      "import importlib.util, os; s=importlib.util.find_spec('libero'); \
-       print(os.path.join(os.path.dirname(s.origin), 'libero'))") && \
-    mkdir -p /home/user_lerobot/.libero && \
-    python -c "\
-from huggingface_hub import snapshot_download; \
-snapshot_download(repo_id='lerobot/libero-assets', repo_type='dataset', \
-                  local_dir='/home/user_lerobot/.libero/assets')" && \
-    printf "assets: /home/user_lerobot/.libero/assets\nbddl_files: ${LIBERO_DIR}/bddl_files\ndatasets: ${LIBERO_DIR}/../datasets\ninit_states: ${LIBERO_DIR}/init_files\n" \
-    > /home/user_lerobot/.libero/config.yaml
-
-# Overlay the PR's source code on top of the nightly image.
-COPY --chown=user_lerobot:user_lerobot . .
-
-CMD ["/bin/bash"]

docker/Dockerfile.benchmark.robomme

@@ -1,56 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Benchmark image for RoboMME integration tests.
-# Extends the nightly GPU image (which has lerobot[all]) with Vulkan system
-# libs for ManiSkill/SAPIEN and the robomme extra. robomme isn't in [all]
-# because mani-skill hard-pins gymnasium==0.29.1 and numpy<2.0.0 which
-# conflict with lerobot's defaults; both are safe at runtime:
-#   - gymnasium 0.29.x has the same 5-tuple step() API as 1.x (since 0.26)
-#   - numpy 1.26.4 is API-compatible with lerobot's actual usage.
-#
-# Build:  docker build -f docker/Dockerfile.benchmark.robomme -t lerobot-benchmark-robomme .
-# Run:    docker run --gpus all --rm lerobot-benchmark-robomme lerobot-eval ...
-
-FROM huggingface/lerobot-gpu:latest
-
-# NVIDIA Container Toolkit: expose Vulkan driver capability for headless rendering.
-ENV NVIDIA_DRIVER_CAPABILITIES=all \
-    VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/nvidia_icd.json
-
-# ManiSkill/SAPIEN's renderer needs Vulkan, which isn't in the base image.
-USER root
-RUN apt-get update \
-    && apt-get install -y --no-install-recommends \
-         libvulkan1 libvulkan-dev mesa-vulkan-drivers \
-    && mkdir -p /usr/share/vulkan/icd.d \
-    && echo '{"file_format_version":"1.0.0","ICD":{"library_path":"libGLX_nvidia.so.0","api_version":"1.3.0"}}' \
-       > /usr/share/vulkan/icd.d/nvidia_icd.json \
-    && apt-get clean && rm -rf /var/lib/apt/lists/*
-USER user_lerobot
-
-# Install smolvla + av-dep via the PR's pyproject, then layer robomme on top
-# with gymnasium/numpy overrides. robomme isn't a pyproject extra because its
-# mani-skill pin conflicts with lerobot's base numpy>=2 (see pyproject.toml).
-COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
-RUN printf 'gymnasium==0.29.1\nnumpy==1.26.4\n' > /tmp/robomme_override.txt \
-    && uv pip install --no-cache --override /tmp/robomme_override.txt \
-         -e ".[smolvla,av-dep]" \
-         "robomme @ git+https://github.com/RoboMME/robomme_benchmark.git@main" \
-    && python -c "import robomme; print('robomme import OK')"
-
-# Overlay the PR's source code on top of the nightly image.
-COPY --chown=user_lerobot:user_lerobot . .
-
-CMD ["/bin/bash"]

docker/Dockerfile.benchmark.robotwin

@@ -1,138 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Benchmark image for RoboTwin 2.0 integration tests.
-# Extends the nightly GPU image with the RoboTwin simulator stack:
-#   sapien/mplib/pytorch3d + NVlabs CuRobo + embodiments.zip + objects.zip
-# (~3.96 GB of assets; background_texture.zip ~11 GB skipped for smoke eval).
-#
-# Build: docker build -f docker/Dockerfile.benchmark.robotwin -t lerobot-benchmark-robotwin .
-# Run:   docker run --gpus all --rm lerobot-benchmark-robotwin \
-#            lerobot-eval --env.type=robotwin --env.task=beat_block_hammer ...
-
-FROM huggingface/lerobot-gpu:latest
-
-ENV NVIDIA_DRIVER_CAPABILITIES=all \
-    VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/nvidia_icd.json \
-    ROBOTWIN_ROOT=/opt/robotwin
-
-# The nightly base is CUDA -base (no compiler, no Vulkan loader). CuRobo's
-# `pip install -e .` runs nvcc, and SAPIEN renders via Vulkan — add both.
-USER root
-# Pinned upstream SHA for reproducible benchmark runs. Bump when we need
-# an upstream fix; don't rely on `main` drift.
-ARG ROBOTWIN_SHA=0aeea2d669c0f8516f4d5785f0aa33ba812c14b4
-RUN apt-get update \
-    && apt-get install -y --no-install-recommends \
-         cuda-nvcc-12-4 cuda-cudart-dev-12-4 \
-         libvulkan1 vulkan-tools \
-    && mkdir -p /usr/share/vulkan/icd.d \
-    && echo '{"file_format_version":"1.0.0","ICD":{"library_path":"libGLX_nvidia.so.0","api_version":"1.3.0"}}' \
-       > /usr/share/vulkan/icd.d/nvidia_icd.json \
-    && git clone https://github.com/RoboTwin-Platform/RoboTwin.git ${ROBOTWIN_ROOT} \
-    && git -C ${ROBOTWIN_ROOT} checkout ${ROBOTWIN_SHA} \
-    && chown -R user_lerobot:user_lerobot ${ROBOTWIN_ROOT} \
-    && apt-get clean && rm -rf /var/lib/apt/lists/*
-USER user_lerobot
-
-# RoboTwin runtime deps (av is already in the base via [av-dep]).
-RUN uv pip install --no-cache \
-        "sapien==3.0.0b1" "mplib==0.2.1" "transforms3d==0.4.2" "trimesh==4.4.3" \
-        "open3d==0.19.0" "imageio==2.34.2" termcolor zarr pydantic h5py
-
-# pytorch3d has no universal wheel; must be built from source (~10 min, cached).
-RUN uv pip install --no-cache --no-build-isolation \
-        "git+https://github.com/facebookresearch/pytorch3d.git@stable"
-
-# CuRobo — NVlabs motion generator; TORCH_CUDA_ARCH_LIST must be set or the
-# build aborts on an empty arch list. RoboTwin's own installer pins v0.7.8,
-# which still exposes the v1 API (`curobo.types.math`) that RoboTwin imports.
-ARG CUROBO_REF=v0.7.8
-RUN cd ${ROBOTWIN_ROOT}/envs \
-    && git clone --branch ${CUROBO_REF} --depth 1 https://github.com/NVlabs/curobo.git \
-    && cd curobo \
-    && TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6;8.9;9.0" \
-       uv pip install -e . --no-build-isolation --no-cache
-
-# Upstream patches (mirror RoboTwin's script/_install.sh).
-# These patches target the exact versions pinned above; re-check when upgrading.
-# mplib==0.2.1: drop a broken `or collide` clause in planner.py.
-#   Safe to remove once mplib > 0.2.1 ships with the fix upstream.
-# sapien==3.0.0b1: fix URDF loader encoding + .srdf extension check.
-#   Safe to remove once sapien > 3.0.0b1 ships with the fix upstream.
-RUN python - <<'EOF'
-import pathlib, re, site
-for d in site.getsitepackages():
-    p = pathlib.Path(d) / "mplib" / "planner.py"
-    if p.exists():
-        p.write_text(re.sub(r"\bor collide\b", "", p.read_text(), count=1))
-        print(f"mplib patch applied: {p}")
-    p = pathlib.Path(d) / "sapien" / "wrapper" / "urdf_loader.py"
-    if p.exists():
-        src = p.read_text().replace(
-            "with open(srdf_path) as f:", 'with open(srdf_path, encoding="utf-8") as f:'
-        ).replace('"srdf"', '".srdf"')
-        p.write_text(src)
-        print(f"sapien patch applied: {p}")
-EOF
-
-# Simulation assets from TianxingChen/RoboTwin2.0: embodiments (~220 MB) +
-# objects (~3.74 GB). background_texture (~11 GB) is intentionally skipped.
-# The dataset is public — no auth token needed.
-RUN python - <<'EOF'
-import os, pathlib, zipfile
-from huggingface_hub import hf_hub_download
-
-assets_dir = pathlib.Path(os.environ["ROBOTWIN_ROOT"]) / "assets"
-assets_dir.mkdir(parents=True, exist_ok=True)
-for fname in ("embodiments.zip", "objects.zip"):
-    local = hf_hub_download(
-        repo_id="TianxingChen/RoboTwin2.0",
-        repo_type="dataset",
-        filename=fname,
-        local_dir=str(assets_dir),
-    )
-    with zipfile.ZipFile(local, "r") as z:
-        z.extractall(str(assets_dir))
-    pathlib.Path(local).unlink()
-EOF
-
-WORKDIR ${ROBOTWIN_ROOT}
-RUN python script/update_embodiment_config_path.py
-
-ENV PYTHONPATH="${ROBOTWIN_ROOT}"
-
-# Fail the image build early if the CuRobo package layout regresses. Importing
-# RoboTwin's planner here is too eager because CuRobo constructs CUDA-backed
-# defaults at import time, while Docker builds don't have access to an NVIDIA
-# driver.
-RUN python - <<'EOF'
-from pathlib import Path
-
-from curobo.types.math import Pose
-
-planner_src = (Path("/opt/robotwin/envs/robot/planner.py")).read_text()
-assert "from curobo.types.math import Pose as CuroboPose" in planner_src
-
-print("CuRobo import OK:", Pose.__name__)
-print("RoboTwin planner import references curobo.types.math")
-EOF
-
-# Return to the lerobot source directory (set by base image) before overlaying.
-WORKDIR /lerobot
-
-# Overlay the PR's source code on top of the nightly image.
-COPY --chown=user_lerobot:user_lerobot . .
-
-CMD ["/bin/bash"]

docker/Dockerfile.benchmark.vlabench

@@ -1,99 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Benchmark image for VLABench integration tests.
-# Extends the nightly GPU image with the PR's source code and VLABench setup.
-#
-# Build:  docker build -f docker/Dockerfile.benchmark.vlabench -t lerobot-benchmark-vlabench .
-# Run:    docker run --gpus all --rm lerobot-benchmark-vlabench lerobot-eval ...
-
-FROM huggingface/lerobot-gpu:latest
-
-# Install VLABench from GitHub (not on PyPI) and pin MuJoCo/dm-control.
-# Shallow-clone without submodule recursion (nested SSH-only submodules fail in CI).
-# Editable install (-e) because VLABench/utils/ has no __init__.py, so
-# find_packages() omits it from wheels; editable mode uses the source tree directly.
-# rrt-algorithms has the same packaging issue (rrt/ dir missing __init__.py).
-# Patch: constant.py calls os.listdir on ~100 asset/obj/meshes/* dirs at import
-# time. Guard the call so missing dirs return [] instead of crashing (in case
-# the asset download is partial).
-#
-# Pinned upstream SHAs for reproducible benchmark runs. Bump when you need
-# an upstream fix; don't rely on `main`/`develop` drift.
-ARG VLABENCH_SHA=cf588fe60c0c7282174fe979f5913170cfe69017
-ARG RRT_ALGORITHMS_SHA=e51d95ee489a225220d6ae2a764c4111f6ba7d85
-RUN git clone https://github.com/OpenMOSS/VLABench.git ~/VLABench && \
-    git -C ~/VLABench checkout ${VLABENCH_SHA} && \
-    git clone https://github.com/motion-planning/rrt-algorithms.git ~/rrt-algorithms && \
-    git -C ~/rrt-algorithms checkout ${RRT_ALGORITHMS_SHA} && \
-    python3 -c "\
-import pathlib; \
-p = pathlib.Path.home() / 'VLABench/VLABench/configs/constant.py'; \
-t = p.read_text(); \
-p.write_text(t.replace( \
-    'subdirs = os.listdir(xml_dir)', \
-    'if not os.path.isdir(xml_dir): return []\n    subdirs = os.listdir(xml_dir)'))" && \
-    uv pip install --no-cache -e ~/VLABench -e ~/rrt-algorithms \
-      mujoco==3.2.2 dm-control==1.0.22 \
-      open3d colorlog scikit-learn openai gdown
-
-# Download VLABench mesh assets. Task configs reference object meshes
-# (obj/meshes/fruit/, containers/basket/, tablewares/plates/, etc.); without
-# them the task builder picks from an empty mesh list and crashes with
-# IndexError at task-build time (random.choice([]) in config_manager.py).
-#
-# Preferred source: an HF Hub mirror. Set VLABENCH_ASSETS_REPO at build time
-# (e.g. --build-arg VLABENCH_ASSETS_REPO=lerobot/vlabench-assets) and we'll
-# snapshot_download the repo into VLABench's assets dir. This is the reliable
-# path for CI — Google Drive frequently returns HTTP 429 ("Too many users have
-# viewed or downloaded this file recently") on shared academic files.
-#
-# After download we *validate* that at least one XML exists under each
-# task-critical subtree and fail the build loudly if not. Silent-empty asset
-# dirs are the #1 cause of VLABench runtime crashes in CI, so we surface them
-# here rather than after a 10-minute eval build.
-#
-# Fallback: VLABench's own gdown-based script. Best-effort only.
-ARG VLABENCH_ASSETS_REPO=""
-RUN ASSETS_DIR="$HOME/VLABench/VLABench/assets" && \
-    if [ -n "${VLABENCH_ASSETS_REPO}" ]; then \
-        echo "Downloading VLABench assets from HF Hub: ${VLABENCH_ASSETS_REPO}" && \
-        uv pip install --no-cache "huggingface_hub[hf_xet]>=0.26" && \
-        python -c "from huggingface_hub import snapshot_download; \
-p = snapshot_download(repo_id='${VLABENCH_ASSETS_REPO}', repo_type='dataset', \
-    local_dir='${ASSETS_DIR}', allow_patterns=['obj/**', 'scenes/**']); \
-print('snapshot_download returned:', p)"; \
-    else \
-        echo "No VLABENCH_ASSETS_REPO set — falling back to gdown" && \
-        python ~/VLABench/scripts/download_assets.py --choice all; \
-    fi && \
-    python -c "\
-from pathlib import Path; \
-import sys; \
-root = Path('${ASSETS_DIR}'); \
-checks = ['obj/meshes/tablewares/plates', 'obj/meshes/containers/basket', 'obj/meshes/fruit', 'obj/meshes/containers/tray']; \
-failed = []; \
-print(f'Validating VLABench assets under {root}'); \
-[print(f'  {c}: {len(list((root/c).rglob(\"*.xml\")))} XMLs') for c in checks]; \
-[failed.append(c) for c in checks if not any((root/c).rglob('*.xml'))]; \
-sys.exit(f'Empty asset dirs (no *.xml): {failed}') if failed else print('All asset dirs populated.')"
-
-# Overlay the PR's source code on top of the nightly image.
-COPY --chown=user_lerobot:user_lerobot . .
-
-# Re-install lerobot editably so the new source (with VLABenchEnv registration
-# and updated obs handling) replaces the stale package baked into the nightly image.
-RUN uv pip install --no-cache --no-deps -e .
-
-CMD ["/bin/bash"]

docs/source/_toctree.yml

@@ -31,12 +31,8 @@
     title: Porting Large Datasets
   - local: using_dataset_tools
     title: Using the Dataset Tools
-  - local: language_and_recipes
-    title: Language Columns and Recipes
-  - local: tools
-    title: Tools
-  - local: annotation_pipeline
-    title: Annotation Pipeline
+  - local: dataset_subtask
+    title: Using Subtasks in the Dataset
   - local: streaming_video_encoding
     title: Streaming Video Encoding
   title: "Datasets"
@@ -65,6 +61,8 @@
     title: SARM
   title: "Reward Models"
 - sections:
+  - local: inference
+    title: Policy Deployment (lerobot-rollout)
   - local: async
     title: Use Async Inference
   - local: rtc
@@ -81,22 +79,10 @@
     title: Adding a New Benchmark
   - local: libero
     title: LIBERO
-  - local: libero_plus
-    title: LIBERO-plus
   - local: metaworld
     title: Meta-World
-  - local: robotwin
-    title: RoboTwin 2.0
-  - local: robocasa
-    title: RoboCasa365
-  - local: robocerebra
-    title: RoboCerebra
-  - local: robomme
-    title: RoboMME
   - local: envhub_isaaclab_arena
     title: NVIDIA IsaacLab Arena Environments
-  - local: vlabench
-    title: VLABench
   title: "Benchmarks"
 - sections:
   - local: introduction_processors

docs/source/annotation_pipeline.mdx

@@ -1,161 +0,0 @@
-# Annotation Pipeline
-
-`lerobot-annotate` populates the two language columns introduced by the
-[Language Columns and Recipes](./language_and_recipes) page —
-`language_persistent` and `language_events` — directly into
-`data/chunk-*/file-*.parquet`. There is no flavor namespace and no sidecar
-file tree: multiple revisions of a dataset mean multiple dataset copies.
-
-## What the pipeline produces
-
-Three modules write into a per-episode staging tree, then a single writer
-rewrites the data shards in place:
-
-| Style / atom                                | Column                | Module   |
-| ------------------------------------------- | --------------------- | -------- |
-| `subtask` (Pi0.7-style "how, not what")     | `language_persistent` | Module 1 |
-| `plan` (initial + refresh on interjection)  | `language_persistent` | Module 1 |
-| `memory` (MEM-style compression)            | `language_persistent` | Module 1 |
-| `interjection`                              | `language_events`     | Module 2 |
-| speech tool-call atom (`style=null`, `say`) | `language_events`     | Module 2 |
-| `vqa` (user / assistant pair)               | `language_events`     | Module 3 |
-
-The writer drops the legacy `subtask_index` column. It does **not** add a
-`tools` column to the parquet — the tool catalog lives at
-`meta/info.json["tools"]` instead (see [Tools](./tools)). After every
-annotation run the pipeline ensures the canonical `say` schema is
-present in that list, preserving any tools the user pre-declared. Chat-
-template consumers read the catalog through
-`LeRobotDatasetMetadata.tools` and pass it to
-`apply_chat_template(messages, tools=meta.tools, ...)`.
-
-If you want to declare additional tools for a dataset before annotation
-runs, edit `meta/info.json["tools"]` directly — the pipeline preserves
-anything already there. Implementations of those tools live under
-`src/lerobot/tools/`; one file per tool, registered via
-`TOOL_REGISTRY`. See the [Tools](./tools) doc for the authoring guide.
-
-## How to run it locally or on SLURM
-
-Install the extra and invoke the console script:
-
-```bash
-uv sync --extra annotations
-uv run lerobot-annotate \
-  --repo_id=imstevenpmwork/super_poulain_draft \
-  --vlm.backend=vllm \
-  --vlm.model_id=Qwen/Qwen3.6-27B-FP8 \
-  --vlm.tensor_parallel_size=2
-```
-
-The pipeline attaches actual camera footage to every Module 1/2/3 prompt
-by default, decoded from the dataset's first `observation.images.*`
-stream. Override with `--vlm.camera_key=observation.images.<name>` to
-pin a specific viewpoint. Datasets with no video tracks fall back to
-text-only prompts automatically.
-
-**Module 1 sees the whole episode as one video block.** Subtask
-decomposition gets a `{"type":"video", "video":[<frames>]}` block
-covering the entire demonstration; Qwen-VL pools temporally on its own
-and decides where to cut. There is no keyframe stride or count knob —
-`--module_1.max_video_frames` (default 32) only caps the frames packed
-into the video block as a model-capacity bound. Module 2 attaches a
-single still frame at the interjection timestamp; Module 3 attaches the
-exact emission frame to each VQA pair.
-
-The executor picks `LocalPipelineExecutor` for small datasets and
-`SlurmPipelineExecutor` for large ones based on
-`--executor.auto_threshold` (default 32 episodes). Force local with
-`--executor.force_local=true`. SLURM jobs honour `--executor.slurm_partition`,
-`--executor.slurm_gpus`, and `--executor.slurm_time`.
-
-## Style-to-recipe consumer mapping
-
-The pipeline produces exactly the styles consumed by
-`src/lerobot/configs/recipes/pi05_hirobot.yaml`:
-
-- `low_level_execution`, `high_level_subtask`, `memory_update` consume
-  `subtask`/`plan`/`memory` from `language_persistent`.
-- `user_interjection_response` consumes `interjection` events plus the
-  paired speech atom (merged into one assistant target turn via
-  `tool_calls_from`) and the same-timestamp `plan` refresh.
-- `ask_vqa` consumes the `(vqa, user)` and `(vqa, assistant)` pairs from
-  `language_events`.
-
-## Why the design is scoped to the canonical recipe
-
-Two things drive the scope:
-
-1. **Persistent state vs exact-event split.** Persistent rows (`subtask`,
-   `plan`, `memory`) broadcast per episode and answer "what state is in
-   force at this frame?". Event rows (`interjection`, `vqa`, speech) only
-   appear on the exact frame whose timestamp matches the emission. The
-   pipeline writes timestamps taken straight from the source parquet — no
-   floating-point recomputation.
-2. **One Qwen-VL pass.** All three modules share a single VLM client
-   (vLLM if available, transformers fallback) so the cost is one model
-   load per dataset, not three.
-
-## Module independence and staged reruns
-
-Each module writes its raw output to
-`<root>/.annotate_staging/episode_{N:06d}/<module>.jsonl`. That makes
-prompt iteration cheap — re-running one module overwrites only its own
-JSONL file before the writer composes the final parquet. Modules can be
-disabled via `--module_1.enabled=false` (and similarly for 2 and 3) to
-test them in isolation.
-
-## Validation/report checks before final write
-
-Before the writer runs, `StagingValidator` checks:
-
-- exact frame-timestamp alignment for every event row;
-- no orphan speech / interjection pairs;
-- `plan` is refreshed at every interjection timestamp;
-- `memory` rows fall on subtask boundaries (warning, not error);
-- VQA assistant `content` parses as JSON in one of the
-  bbox / keypoint / count / attribute / spatial shapes;
-- every row routes to the column dictated by `column_for_style(style)`.
-
-Errors abort the writer (`--skip_validation=true` overrides for debugging).
-
-## Paper inspirations per module
-
-- **Module 1 — subtasks.** Hi Robot ([Shi 2025](https://arxiv.org/abs/2502.19417))
-  atom granularity ("pick up one piece of lettuce", "place bowl to box");
-  Pi0.7 ([Physical Intelligence 2025](https://pi.website/pi07)) "how, not
-  what" detail.
-- **Module 1 — memory.** MEM ([Torne 2026](https://arxiv.org/abs/2603.03596))
-  compression directive: keep only minimal relevant information; functional
-  outcomes preserved, specific attributes dropped.
-- **Module 2 — interjections.** Hi Robot scenario taxonomy: negative task,
-  situated correction, specific constraint, preference. Speech is a
-  tool-call-only atom (`tool_calls=[{type:function, function:{name:"say",
-arguments:{text:...}}}]`).
-- **Module 3 — VQA.** ECoT ([Zawalski 2024](https://arxiv.org/abs/2407.08693))
-  grounded features (bounding boxes in pixel `[x_min, y_min, x_max, y_max]`,
-  keypoints) and Steerable Policies' multi-abstraction grounding.
-
-Future maintainers should adjust the prompt templates in
-`src/lerobot/annotations/steerable_pipeline/prompts/` against these
-references rather than rewriting from scratch.
-
-## Compute and list-size estimates
-
-Per episode, the pipeline issues O(`max_steps`) Module 1 calls,
-O(`max_interjections_per_episode`) Module 2 calls, and
-O(`vqa_emission_hz × episode_seconds`) Module 3 calls. With defaults
-(8 subtasks, 1 interjection, 1 Hz × 3 pairs) and 30-second episodes, that
-is ~50 VLM calls per episode. `language_persistent` per episode is ~10s of
-KB at most (parquet dictionary-encodes one entry per episode);
-`language_events` is empty on most frames and is bounded by the number of
-emissions, not `num_frames × num_emissions`.
-
-## Reproducibility via seed and prompt hashes
-
-`--seed` (default 1729) feeds the per-episode RNGs that select interjection
-timestamps and VQA question types. Combined with the deterministic prompt
-templates checked into `prompts/`, two runs at the same seed against the
-same dataset and the same model checkpoint produce byte-identical staging
-artifacts. Prompt edits are recorded by file hash; future tooling can pin
-expected `(seed, prompt_hash)` pairs into the dataset card.

docs/source/dataset_subtask.mdx

@@ -0,0 +1,277 @@
+# Using Subtasks in LeRobot Datasets
+
+Subtask support in robotics datasets has proven effective in improving robot reasoning and understanding. Subtasks are particularly useful for:
+
+- **Hierarchical policies**: Building policies that include subtask predictions to visualize robot reasoning in real time
+- **Reward modeling**: Helping reward models understand task progression (e.g., SARM-style stage-aware reward models)
+- **Task decomposition**: Breaking down complex manipulation tasks into atomic, interpretable steps
+
+LeRobotDataset now supports subtasks as part of its dataset structure, alongside tasks.
+
+## What are Subtasks?
+
+While a **task** describes the overall goal (e.g., "Pick up the apple and place it in the basket"), **subtasks** break down the execution into finer-grained steps:
+
+1. "Approach the apple"
+2. "Grasp the apple"
+3. "Lift the apple"
+4. "Move to basket"
+5. "Release the apple"
+
+Each frame in the dataset can be annotated with its corresponding subtask, enabling models to learn and predict these intermediate stages.
+
+<img
+  src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/subtask-asset.png"
+  alt="An overview of subtask annotation showing how frames are labeled with intermediate subtask stages"
+  width="80%"
+/>
+
+<p>
+  <em>Figure: Overview of subtask annotation.</em>
+</p>
+
+**Reference:** _Subtask-learning based for robot self-assembly in flexible collaborative assembly in manufacturing_, Original Article, Published: 19 April 2022.
+
+## Dataset Structure
+
+Subtask information is stored in the dataset metadata:
+
+```
+my-dataset/
+├── data/
+│   └── ...
+├── meta/
+│   ├── info.json
+│   ├── stats.json
+│   ├── tasks.parquet
+│   ├── subtasks.parquet      # Subtask index → subtask string mapping
+│   └── episodes/
+│       └── ...
+└── videos/
+    └── ...
+```
+
+### Subtasks Parquet File
+
+The `meta/subtasks.parquet` file maps subtask indices to their natural language descriptions:
+
+| subtask_index | subtask (index column) |
+| ------------- | ---------------------- |
+| 0             | "Approach the apple"   |
+| 1             | "Grasp the apple"      |
+| 2             | "Lift the apple"       |
+| ...           | ...                    |
+
+### Frame-Level Annotations
+
+Each frame in the dataset can include a `subtask_index` field that references the subtasks parquet file:
+
+```python
+# Example frame data in the parquet file
+{
+    "index": 42,
+    "timestamp": 1.4,
+    "episode_index": 0,
+    "task_index": 0,
+    "subtask_index": 2,  # References "Lift the apple"
+    "observation.state": [...],
+    "action": [...],
+}
+```
+
+## Annotating Datasets with Subtasks
+
+We provide a HuggingFace Space for easily annotating any LeRobotDataset with subtasks:
+
+**[https://huggingface.co/spaces/lerobot/annotate](https://huggingface.co/spaces/lerobot/annotate)**
+
+After completing your annotation:
+
+1. Click "Push to Hub" to upload your annotated dataset
+2. You can also run the annotation space locally by following the instructions at [github.com/huggingface/lerobot-annotate](https://github.com/huggingface/lerobot-annotate)
+
+## Loading Datasets with Subtasks
+
+When you load a dataset with subtask annotations, the subtask information is automatically available:
+
+```python
+from lerobot.datasets import LeRobotDataset
+
+# Load a dataset with subtask annotations
+dataset = LeRobotDataset("jadechoghari/collect-fruit-annotated")
+
+# Access a sample
+sample = dataset[100]
+
+# The sample includes both task and subtask information
+print(sample["task"])        # "Collect the fruit"
+print(sample["subtask"])     # "Grasp the apple"
+print(sample["task_index"])  # tensor(0)
+print(sample["subtask_index"])  # tensor(2)
+```
+
+### Checking for Subtask Support
+
+You can check if a dataset has subtask annotations:
+
+```python
+# Check if subtasks are available
+has_subtasks = (
+    "subtask_index" in dataset.features
+    and dataset.meta.subtasks is not None
+)
+
+if has_subtasks:
+    print(f"Dataset has {len(dataset.meta.subtasks)} unique subtasks")
+    print("Subtasks:", list(dataset.meta.subtasks.index))
+```
+
+## Using Subtasks for Training
+
+### With the Tokenizer Processor
+
+The `TokenizerProcessor` automatically handles subtask tokenization for Vision-Language Action (VLA) models:
+
+```python
+from lerobot.processor import TokenizerProcessorStep
+
+# Create a tokenizer processor step
+tokenizer_processor = TokenizerProcessorStep(
+    tokenizer_name_or_path="google/paligemma-3b-pt-224",
+    padding="max_length",
+    max_length=64,
+)
+
+# The processor will automatically tokenize subtasks if present in the batch
+# and add them to the observation under:
+# - "observation.subtask.tokens"
+# - "observation.subtask.attention_mask"
+```
+
+When subtasks are available in the batch, the tokenizer processor adds:
+
+- `observation.subtask.tokens`: Tokenized subtask text
+- `observation.subtask.attention_mask`: Attention mask for the subtask tokens
+
+### DataLoader with Subtasks
+
+```python
+import torch
+from lerobot.datasets import LeRobotDataset
+
+dataset = LeRobotDataset("jadechoghari/collect-fruit-annotated")
+
+dataloader = torch.utils.data.DataLoader(
+    dataset,
+    batch_size=16,
+    shuffle=True,
+)
+
+for batch in dataloader:
+    # Access subtask information in the batch
+    subtasks = batch["subtask"]  # List of subtask strings
+    subtask_indices = batch["subtask_index"]  # Tensor of subtask indices
+
+    # Use for training hierarchical policies or reward models
+    print(f"Batch subtasks: {set(subtasks)}")
+```
+
+## Example Datasets with Subtask Annotations
+
+Try loading a dataset with subtask annotations:
+
+```python
+from lerobot.datasets import LeRobotDataset
+
+# Example dataset with subtask annotations
+dataset = LeRobotDataset("jadechoghari/collect-fruit-annotated")
+
+# Explore the subtasks
+print("Available subtasks:")
+for subtask_name in dataset.meta.subtasks.index:
+    print(f"  - {subtask_name}")
+
+# Get subtask distribution
+subtask_counts = {}
+for i in range(len(dataset)):
+    sample = dataset[i]
+    subtask = sample["subtask"]
+    subtask_counts[subtask] = subtask_counts.get(subtask, 0) + 1
+
+print("\nSubtask distribution:")
+for subtask, count in sorted(subtask_counts.items(), key=lambda x: -x[1]):
+    print(f"  {subtask}: {count} frames")
+```
+
+## Use Cases
+
+### 1. Hierarchical Policy Training
+
+Train policies that predict both actions and current subtask:
+
+```python
+class HierarchicalPolicy(nn.Module):
+    def __init__(self, num_subtasks):
+        super().__init__()
+        self.action_head = nn.Linear(hidden_dim, action_dim)
+        self.subtask_head = nn.Linear(hidden_dim, num_subtasks)
+
+    def forward(self, observations):
+        features = self.encoder(observations)
+        actions = self.action_head(features)
+        subtask_logits = self.subtask_head(features)
+        return actions, subtask_logits
+```
+
+### 2. Stage-Aware Reward Modeling (SARM)
+
+Build reward models that understand task progression:
+
+```python
+# SARM predicts:
+# - Stage: Which subtask is being executed (discrete)
+# - Progress: How far along the subtask (continuous 0-1)
+
+class SARMRewardModel(nn.Module):
+    def forward(self, observations):
+        features = self.encoder(observations)
+        stage_logits = self.stage_classifier(features)
+        progress = self.progress_regressor(features)
+        return stage_logits, progress
+```
+
+### 3. Progress Visualization
+
+Monitor robot execution by tracking subtask progression:
+
+```python
+def visualize_execution(model, observations):
+    for t, obs in enumerate(observations):
+        action, subtask_logits = model(obs)
+        predicted_subtask = subtask_names[subtask_logits.argmax()]
+        print(f"t={t}: Executing '{predicted_subtask}'")
+```
+
+## API Reference
+
+### LeRobotDataset Properties
+
+| Property                    | Type                   | Description                                |
+| --------------------------- | ---------------------- | ------------------------------------------ |
+| `meta.subtasks`             | `pd.DataFrame \| None` | DataFrame mapping subtask names to indices |
+| `features["subtask_index"]` | `dict`                 | Feature spec for subtask_index if present  |
+
+### Sample Keys
+
+When subtasks are available, each sample includes:
+
+| Key             | Type           | Description                          |
+| --------------- | -------------- | ------------------------------------ |
+| `subtask_index` | `torch.Tensor` | Integer index of the current subtask |
+| `subtask`       | `str`          | Natural language subtask description |
+
+## Related Resources
+
+- [SARM Paper](https://arxiv.org/pdf/2509.25358) - Stage-Aware Reward Modeling for Long Horizon Robot Manipulation
+- [LeRobot Annotate Space](https://huggingface.co/spaces/lerobot/annotate) - Interactive annotation tool
+- [LeRobotDataset v3.0](./lerobot-dataset-v3) - Dataset format documentation

docs/source/hil_data_collection.mdx

@@ -50,30 +50,30 @@ This process can be repeated iteratively: deploy, collect, fine-tune, repeat. Ea
 
 ### Teleoperator Requirements
 
-The `examples/hil` HIL scripts require **teleoperators with active motors** that can:
+The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with active motors** that can:
 
 - Enable/disable torque programmatically
 - Move to target positions (to mirror the robot state when pausing)
 
-**Compatible teleoperators in the current `examples/hil` scripts:**
+**Compatible teleoperators:**
 
 - `openarm_mini` - OpenArm Mini
 - `so_leader` - SO100 / SO101 leader arm
 
 > [!IMPORTANT]
-> The provided `examples/hil` commands default to `bi_openarm_follower` + `openarm_mini`.
+> The provided commands default to `bi_openarm_follower` + `openarm_mini`.
 > `so_follower` + `so_leader` configs are also registered and can be used via CLI flags.
 
 ---
 
 ## Script
 
-A single script handles both synchronous and RTC-based inference. Toggle RTC with `--rtc.enabled=true`:
+Use `lerobot-rollout` with `--strategy.type=dagger` for HIL data collection. Select the inference backend with `--inference.type=sync|rtc`:
 
-| Mode                     | Flag                 | Models                |
-| ------------------------ | -------------------- | --------------------- |
-| Standard (default)       | _(no flag needed)_   | ACT, Diffusion Policy |
-| Real-Time Chunking (RTC) | `--rtc.enabled=true` | Pi0, Pi0.5, SmolVLA   |
+| Mode                     | Flag                   | Models                |
+| ------------------------ | ---------------------- | --------------------- |
+| Standard (default)       | _(no flag needed)_     | ACT, Diffusion Policy |
+| Real-Time Chunking (RTC) | `--inference.type=rtc` | Pi0, Pi0.5, SmolVLA   |
 
 ---
 
@@ -97,7 +97,7 @@ python src/lerobot/scripts/lerobot_train.py \
 **Standard inference (ACT, Diffusion Policy):**
 
 ```bash
-python examples/hil/hil_data_collection.py \
+lerobot-rollout --strategy.type=dagger \
     --robot.type=bi_openarm_follower \
     --robot.left_arm_config.port=can1 \
     --robot.left_arm_config.side=left \
@@ -111,8 +111,7 @@ python examples/hil/hil_data_collection.py \
     --dataset.repo_id=your-username/hil-dataset \
     --dataset.single_task="Fold the T-shirt properly" \
     --dataset.fps=30 \
-    --dataset.episode_time_s=1000 \
-    --dataset.num_episodes=50 \
+    --strategy.num_episodes=50 \
     --interpolation_multiplier=2
 ```
 
@@ -121,11 +120,11 @@ python examples/hil/hil_data_collection.py \
 For models with high inference latency, enable RTC for smooth execution:
 
 ```bash
-python examples/hil/hil_data_collection.py \
-    --rtc.enabled=true \
-    --rtc.execution_horizon=20 \
-    --rtc.max_guidance_weight=5.0 \
-    --rtc.prefix_attention_schedule=LINEAR \
+lerobot-rollout --strategy.type=dagger \
+    --inference.type=rtc \
+    --inference.rtc.execution_horizon=20 \
+    --inference.rtc.max_guidance_weight=5.0 \
+    --inference.rtc.prefix_attention_schedule=LINEAR \
     --robot.type=bi_openarm_follower \
     --robot.left_arm_config.port=can1 \
     --robot.left_arm_config.side=left \
@@ -139,8 +138,7 @@ python examples/hil/hil_data_collection.py \
     --dataset.repo_id=your-username/hil-rtc-dataset \
     --dataset.single_task="Fold the T-shirt properly" \
     --dataset.fps=30 \
-    --dataset.episode_time_s=1000 \
-    --dataset.num_episodes=50 \
+    --strategy.num_episodes=50 \
     --interpolation_multiplier=3
 ```
 
@@ -235,7 +233,7 @@ This HIL data collection approach builds on ideas from interactive imitation lea
 
 - **HG-DAgger** (Kelly et al., 2019) made this practical for robotics: a human expert monitors the robot and only intervenes when needed, rather than labeling every state. The gating between autonomous and human control is exactly the pause → takeover → return-to-policy loop used in the scripts here.
 
-- **RaC** (Hu et al., 2025) scales this loop to long-horizon tasks by explicitly decomposing interventions into **recovery** (teleoperating back to a good state) and **correction** (demonstrating the right behavior from there). This decomposition is the protocol followed by the HIL scripts in `examples/hil`.
+- **RaC** (Hu et al., 2025) scales this loop to long-horizon tasks by explicitly decomposing interventions into **recovery** (teleoperating back to a good state) and **correction** (demonstrating the right behavior from there). This decomposition is the protocol followed by the DAgger strategy in `lerobot-rollout`.
 
 - **π0.6/RECAP** (Physical Intelligence, 2025) applies the same iterative collect-and-finetune loop at scale with VLA models, showing that even large pretrained policies benefit substantially from targeted human corrections on their own failure modes. π0.6 is trained using RECAP.
 

docs/source/il_robots.mdx

@@ -509,121 +509,42 @@ hf upload ${HF_USER}/act_so101_test${CKPT} \
 
 ## Run inference and evaluate your policy
 
-You can use the `record` script from [`lerobot-record`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/scripts/lerobot_record.py) with a policy checkpoint as input, to run inference and evaluate your policy. For instance, run this command or API example to run inference and record 10 evaluation episodes:
+Use `lerobot-rollout` to deploy a trained policy on your robot. You can choose different strategies depending on your needs:
 
 <hfoptions id="eval">
-<hfoption id="Command">
+<hfoption id="Base mode (no recording)">
 ```bash
-lerobot-record  \
+lerobot-rollout \
+  --strategy.type=base \
+  --policy.path=${HF_USER}/my_policy \
   --robot.type=so100_follower \
   --robot.port=/dev/ttyACM1 \
   --robot.cameras="{ up: {type: opencv, index_or_path: /dev/video10, width: 640, height: 480, fps: 30}, side: {type: intelrealsense, serial_number_or_name: 233522074606, width: 640, height: 480, fps: 30}}" \
-  --robot.id=my_awesome_follower_arm \
-  --display_data=false \
-  --dataset.repo_id=${HF_USER}/eval_so100 \
-  --dataset.single_task="Put lego brick into the transparent box" \
-  --dataset.streaming_encoding=true \
-  --dataset.encoder_threads=2 \
-  # --dataset.vcodec=auto \
-  # <- Teleop optional if you want to teleoperate in between episodes \
-  # --teleop.type=so100_leader \
-  # --teleop.port=/dev/ttyACM0 \
-  # --teleop.id=my_awesome_leader_arm \
-  --policy.path=${HF_USER}/my_policy
+  --task="Put lego brick into the transparent box" \
+  --duration=60
 ```
 </hfoption>
-<hfoption id="API example">
-
-<!-- prettier-ignore-start -->
-```python
-from lerobot.cameras.opencv import OpenCVCameraConfig
-from lerobot.datasets import LeRobotDataset
-from lerobot.utils.feature_utils import hw_to_dataset_features
-from lerobot.policies.act import ACTPolicy
-from lerobot.policies import make_pre_post_processors
-from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
-from lerobot.scripts.lerobot_record import record_loop
-from lerobot.common.control_utils import init_keyboard_listener
-from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
-
-
-NUM_EPISODES = 5
-FPS = 30
-EPISODE_TIME_SEC = 60
-TASK_DESCRIPTION = "My task description"
-HF_MODEL_ID = "<hf_username>/<model_repo_id>"
-HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>"
-
-# Create the robot configuration
-camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
-robot_config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", cameras=camera_config
-)
-
-# Initialize the robot
-robot = SO100Follower(robot_config)
-
-# Initialize the policy
-policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
-
-# Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, "action")
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
-dataset_features = {**action_features, **obs_features}
-
-# Create the dataset
-dataset = LeRobotDataset.create(
-    repo_id=HF_DATASET_ID,
-    fps=FPS,
-    features=dataset_features,
-    robot_type=robot.name,
-    use_videos=True,
-    image_writer_threads=4,
-)
-
-# Initialize the keyboard listener and rerun visualization
-_, events = init_keyboard_listener()
-init_rerun(session_name="recording")
-
-# Connect the robot
-robot.connect()
-
-preprocessor, postprocessor = make_pre_post_processors(
-    policy_cfg=policy,
-    pretrained_path=HF_MODEL_ID,
-    dataset_stats=dataset.meta.stats,
-)
-
-for episode_idx in range(NUM_EPISODES):
-    log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
-
-    # Run the policy inference loop
-    record_loop(
-        robot=robot,
-        events=events,
-        fps=FPS,
-        policy=policy,
-        preprocessor=preprocessor,
-        postprocessor=postprocessor,
-        dataset=dataset,
-        control_time_s=EPISODE_TIME_SEC,
-        single_task=TASK_DESCRIPTION,
-        display_data=True,
-    )
-
-    dataset.save_episode()
-
-# Clean up
-robot.disconnect()
-dataset.push_to_hub()
+<hfoption id="Sentry mode (with recording)">
+```bash
+lerobot-rollout \
+  --strategy.type=sentry \
+  --strategy.upload_every_n_episodes=5 \
+  --policy.path=${HF_USER}/my_policy \
+  --robot.type=so100_follower \
+  --robot.port=/dev/ttyACM1 \
+  --robot.cameras="{ up: {type: opencv, index_or_path: /dev/video10, width: 640, height: 480, fps: 30}, side: {type: intelrealsense, serial_number_or_name: 233522074606, width: 640, height: 480, fps: 30}}" \
+  --dataset.repo_id=${HF_USER}/eval_so100 \
+  --dataset.single_task="Put lego brick into the transparent box" \
+  --duration=600
 ```
-<!-- prettier-ignore-end -->
-
 </hfoption>
 </hfoptions>
 
-As you can see, it's almost the same command as previously used to record your training dataset. Two things changed:
+The `--strategy.type` flag selects the execution mode:
 
-1. There is an additional `--control.policy.path` argument which indicates the path to your policy checkpoint with (e.g. `outputs/train/eval_act_so101_test/checkpoints/last/pretrained_model`). You can also use the model repository if you uploaded a model checkpoint to the hub (e.g. `${HF_USER}/act_so101_test`).
-2. The name of dataset begins by `eval` to reflect that you are running inference (e.g. `${HF_USER}/eval_act_so101_test`).
+- `base`: Autonomous rollout with no data recording (useful for quick evaluation)
+- `sentry`: Continuous recording with auto-upload (useful for large-scale evaluation)
+- `highlight`: Ring buffer recording with keystroke save (useful for capturing interesting events)
+- `dagger`: Human-in-the-loop data collection (see [HIL Data Collection](./hil_data_collection))
+
+All strategies support `--inference.type=rtc` for smooth execution with slow VLA models (Pi0, Pi0.5, SmolVLA).

docs/source/inference.mdx

@@ -0,0 +1,261 @@
+# Policy Deployment (lerobot-rollout)
+
+`lerobot-rollout` is the single CLI for deploying trained policies on real robots. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording and human-in-the-loop data collection.
+
+## Quick Start
+
+No extra dependencies are needed beyond your robot and policy extras.
+
+```bash
+lerobot-rollout \
+    --strategy.type=base \
+    --policy.path=lerobot/act_koch_real \
+    --robot.type=koch_follower \
+    --robot.port=/dev/ttyACM0 \
+    --task="pick up cube" \
+    --duration=30
+```
+
+This runs the policy for 30 seconds with no recording.
+
+---
+
+## Strategies
+
+Select a strategy with `--strategy.type=<name>`. Each strategy defines a different control loop with its own recording and interaction semantics.
+
+### Base (`--strategy.type=base`)
+
+Autonomous policy execution with no data recording. Use this for quick evaluation, demos, or when you only need to observe the robot.
+
+```bash
+lerobot-rollout \
+    --strategy.type=base \
+    --policy.path=${HF_USER}/my_policy \
+    --robot.type=so100_follower \
+    --robot.port=/dev/ttyACM0 \
+    --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
+    --task="Put lego brick into the box" \
+    --duration=60
+```
+
+| Flag             | Description                                            |
+| ---------------- | ------------------------------------------------------ |
+| `--duration`     | Run time in seconds (0 = infinite)                     |
+| `--task`         | Task description passed to the policy                  |
+| `--display_data` | Stream observations/actions to Rerun for visualization |
+
+### Sentry (`--strategy.type=sentry`)
+
+Continuous autonomous recording with periodic upload to the Hugging Face Hub. Episode boundaries are auto-computed from camera resolution and FPS so each saved episode produces a complete video file, keeping uploads efficient.
+
+Policy state (hidden state, RTC queue) persists across episode boundaries: the robot does not reset between episodes.
+
+```bash
+lerobot-rollout \
+    --strategy.type=sentry \
+    --strategy.upload_every_n_episodes=5 \
+    --policy.path=${HF_USER}/my_policy \
+    --robot.type=so100_follower \
+    --robot.port=/dev/ttyACM0 \
+    --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
+    --dataset.repo_id=${HF_USER}/eval_data \
+    --dataset.single_task="Put lego brick into the box" \
+    --duration=3600
+```
+
+| Flag                                   | Description                                                 |
+| -------------------------------------- | ----------------------------------------------------------- |
+| `--strategy.upload_every_n_episodes`   | Push to Hub every N episodes (default: 5)                   |
+| `--strategy.target_video_file_size_mb` | Target video file size for episode rotation (default: auto) |
+| `--dataset.repo_id`                    | **Required.** Hub repository for the recorded dataset       |
+| `--dataset.push_to_hub`                | Whether to push to Hub on teardown (default: true)          |
+
+### Highlight (`--strategy.type=highlight`)
+
+Autonomous rollout with on-demand recording via a memory-bounded ring buffer. The robot runs continuously while the buffer captures the last N seconds of telemetry. Press the save key to flush the buffer and start live recording; press it again to save the episode.
+
+```bash
+lerobot-rollout \
+    --strategy.type=highlight \
+    --strategy.ring_buffer_seconds=30 \
+    --strategy.save_key=s \
+    --strategy.push_key=h \
+    --policy.path=${HF_USER}/my_policy \
+    --robot.type=koch_follower \
+    --robot.port=/dev/ttyACM0 \
+    --dataset.repo_id=${HF_USER}/highlight_data \
+    --dataset.single_task="Pick up the red cube"
+```
+
+**Keyboard controls:**
+
+| Key                | Action                                                   |
+| ------------------ | -------------------------------------------------------- |
+| `s` (configurable) | Start recording (flushes buffer) / stop and save episode |
+| `h` (configurable) | Push dataset to Hub                                      |
+| `ESC`              | Stop the session                                         |
+
+| Flag                                   | Description                                    |
+| -------------------------------------- | ---------------------------------------------- |
+| `--strategy.ring_buffer_seconds`       | Duration of buffered telemetry (default: 30)   |
+| `--strategy.ring_buffer_max_memory_mb` | Memory cap for the ring buffer (default: 2048) |
+| `--strategy.save_key`                  | Key to toggle recording (default: `s`)         |
+| `--strategy.push_key`                  | Key to push to Hub (default: `h`)              |
+
+### DAgger (`--strategy.type=dagger`)
+
+Human-in-the-loop data collection. Alternates between autonomous policy execution and human intervention via a teleoperator. Intervention frames are tagged with `intervention=True`. Requires a teleoperator (`--teleop.type`).
+
+See the [Human-In-the-Loop Data Collection](./hil_data_collection) guide for a detailed walkthrough.
+
+**Corrections-only mode** (default): Only human correction windows are recorded. Each correction becomes one episode.
+
+```bash
+lerobot-rollout \
+    --strategy.type=dagger \
+    --strategy.num_episodes=20 \
+    --policy.path=outputs/pretrain/checkpoints/last/pretrained_model \
+    --robot.type=bi_openarm_follower \
+    --teleop.type=openarm_mini \
+    --dataset.repo_id=${HF_USER}/hil_data \
+    --dataset.single_task="Fold the T-shirt"
+```
+
+**Continuous recording mode** (`--strategy.record_autonomous=true`): Both autonomous and correction frames are recorded with time-based episode rotation (same as Sentry).
+
+```bash
+lerobot-rollout \
+    --strategy.type=dagger \
+    --strategy.record_autonomous=true \
+    --strategy.num_episodes=50 \
+    --policy.path=${HF_USER}/my_policy \
+    --robot.type=so100_follower \
+    --robot.port=/dev/ttyACM0 \
+    --teleop.type=so101_leader \
+    --teleop.port=/dev/ttyACM1 \
+    --dataset.repo_id=${HF_USER}/dagger_data \
+    --dataset.single_task="Grasp the block"
+```
+
+**Keyboard controls** (default input device):
+
+| Key     | Action                                      |
+| ------- | ------------------------------------------- |
+| `Space` | Pause / resume policy execution             |
+| `Tab`   | Start / stop human correction               |
+| `Enter` | Push dataset to Hub (corrections-only mode) |
+| `ESC`   | Stop the session                            |
+
+Foot pedal input is also supported via `--strategy.input_device=pedal`. Configure pedal codes with `--strategy.pedal.*` flags.
+
+| Flag                                 | Description                                             |
+| ------------------------------------ | ------------------------------------------------------- |
+| `--strategy.num_episodes`            | Number of correction episodes to record (default: 10)   |
+| `--strategy.record_autonomous`       | Record autonomous frames too (default: false)           |
+| `--strategy.upload_every_n_episodes` | Push to Hub every N episodes (default: 5)               |
+| `--strategy.input_device`            | Input device: `keyboard` or `pedal` (default: keyboard) |
+| `--teleop.type`                      | **Required.** Teleoperator type                         |
+
+---
+
+## Inference Backends
+
+Select a backend with `--inference.type=<name>`. All strategies work with both backends.
+
+### Sync (default)
+
+One policy call per control tick. The main loop blocks until the action is computed.
+
+Works with all policies. No extra flags needed.
+
+### Real-Time Chunking (`--inference.type=rtc`)
+
+A background thread produces action chunks asynchronously. The main control loop polls for the next ready action while the policy computes the next chunk in parallel.
+
+Use RTC with large, slow VLA models (Pi0, Pi0.5, SmolVLA) for smooth, continuous motion despite high inference latency.
+
+```bash
+lerobot-rollout \
+    --strategy.type=base \
+    --inference.type=rtc \
+    --inference.rtc.execution_horizon=10 \
+    --inference.rtc.max_guidance_weight=10.0 \
+    --policy.path=${HF_USER}/pi0_policy \
+    --robot.type=so100_follower \
+    --robot.port=/dev/ttyACM0 \
+    --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
+    --task="Pick up the cube" \
+    --duration=60 \
+    --device=cuda
+```
+
+| Flag                                        | Description                                                    |
+| ------------------------------------------- | -------------------------------------------------------------- |
+| `--inference.rtc.execution_horizon`         | Steps to blend with previous chunk (default: varies by policy) |
+| `--inference.rtc.max_guidance_weight`       | Consistency enforcement strength (default: varies by policy)   |
+| `--inference.rtc.prefix_attention_schedule` | Blend schedule: `LINEAR`, `EXP`, `ONES`, `ZEROS`               |
+| `--inference.queue_threshold`               | Max queue size before backpressure (default: 30)               |
+
+See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters.
+
+---
+
+## Common Flags
+
+| Flag                              | Description                                                       | Default |
+| --------------------------------- | ----------------------------------------------------------------- | ------- |
+| `--policy.path`                   | **Required.** HF Hub model ID or local checkpoint path            | --      |
+| `--robot.type`                    | **Required.** Robot type (e.g. `so100_follower`, `koch_follower`) | --      |
+| `--robot.port`                    | Serial port for the robot                                         | --      |
+| `--robot.cameras`                 | Camera configuration (JSON dict)                                  | --      |
+| `--fps`                           | Control loop frequency                                            | 30      |
+| `--duration`                      | Run time in seconds (0 = infinite)                                | 0       |
+| `--device`                        | Torch device (`cpu`, `cuda`, `mps`)                               | auto    |
+| `--task`                          | Task description (used when no dataset is provided)               | --      |
+| `--display_data`                  | Stream telemetry to Rerun visualization                           | false   |
+| `--display_ip` / `--display_port` | Remote Rerun server address                                       | --      |
+| `--interpolation_multiplier`      | Action interpolation factor                                       | 1       |
+| `--use_torch_compile`             | Enable `torch.compile` for inference                              | false   |
+| `--resume`                        | Resume a previous recording session                               | false   |
+| `--play_sounds`                   | Vocal synthesis for events                                        | true    |
+
+---
+
+## Programmatic Usage
+
+For custom deployments (e.g. with kinematics processors), use the rollout module API directly:
+
+```python
+from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
+from lerobot.rollout.inference import SyncInferenceConfig
+from lerobot.rollout.strategies import BaseStrategy
+from lerobot.utils.process import ProcessSignalHandler
+
+cfg = RolloutConfig(
+    robot=my_robot_config,
+    policy=my_policy_config,
+    strategy=BaseStrategyConfig(),
+    inference=SyncInferenceConfig(),
+    fps=30,
+    duration=60,
+    task="my task",
+)
+
+signal_handler = ProcessSignalHandler(use_threads=True)
+ctx = build_rollout_context(
+    cfg,
+    signal_handler.shutdown_event,
+    robot_action_processor=my_custom_action_processor,       # optional
+    robot_observation_processor=my_custom_obs_processor,     # optional
+)
+
+strategy = BaseStrategy(cfg.strategy)
+try:
+    strategy.setup(ctx)
+    strategy.run(ctx)
+finally:
+    strategy.teardown(ctx)
+```
+
+See `examples/so100_to_so100_EE/rollout.py` and `examples/phone_to_so100/rollout.py` for full examples with kinematics processors.

docs/source/language_and_recipes.mdx

@@ -1,109 +0,0 @@
-# Language columns and recipes
-
-LeRobot stores reusable language annotations directly next to frame data in `data/chunk-*/file-*.parquet`.
-The two optional columns are:
-
-- `language_persistent`: a list of rows broadcast across every frame in an episode for state that remains active, such as `subtask`, `plan`, and `memory`.
-- `language_events`: a list of rows only on the exact frame where an event was emitted, such as `interjection`, `vqa`, and speech tool calls.
-
-Both columns share the same row shape (event rows omit `timestamp` because the
-frame the row sits on already provides it):
-
-```text
-role: string
-content: string | null
-style: string | null
-timestamp: float64        # persistent rows only
-camera: string | null     # observation.images.* feature key, view-dependent rows only
-tool_calls: list[Json] | null
-```
-
-The `camera` field tags rows whose `content` is grounded in a specific camera
-view. Rows of view-dependent styles (`vqa`, and the reserved `motion` /
-`trace`) MUST set `camera` to the matching `observation.images.*` feature key.
-Rows of every other style MUST leave `camera` as `null`. Pipeline writers and
-the validator enforce this via `validate_camera_field(style, camera)`.
-
-`meta/tasks.parquet` remains the canonical source for the task. The special `${task}` recipe binding always reads that task string and does not depend on language annotations.
-
-## Architecture
-
-The language stack has three layers:
-
-1. `lerobot.datasets.language` defines the schema, style registry, and `column_for_style`.
-2. `lerobot.datasets.language_render` resolves rows and renders messages.
-3. `RenderMessagesStep` turns dataset samples into `messages`, `message_streams`, and `target_message_indices`.
-
-`LeRobotDataset` stays recipe-agnostic. It passes `language_persistent` and `language_events` through when present, and unannotated datasets keep their existing behavior.
-
-## Temporal semantics
-
-Persistent styles are active after emission until replaced:
-
-- `active_at(t, style=subtask)`
-- `nth_prev(style=memory, offset=1)`
-- `nth_next(style=subtask, offset=1)`
-
-Event styles only exist on their exact timestamp:
-
-- `emitted_at(t, style=interjection)`
-- `emitted_at(t, style=vqa, role=user, camera=observation.images.top)`
-- `emitted_at(t, role=assistant, tool_name=say)`
-
-Exact event matching has no tolerance window, so writers must stamp event rows with frame timestamps from the parquet data.
-
-## View-dependent resolution
-
-For view-dependent styles (`vqa`, `motion`, `trace`), the resolver gains a
-`camera=` filter parallel to `role=` and `tool_name=`. Datasets with multiple
-cameras typically emit one (`vqa`, `user`) + (`vqa`, `assistant`) pair per
-camera at the same timestamp; without `camera=`, those resolvers see two
-matches and raise an ambiguity error. Recipes consume each camera through its
-own binding plus a matching image block, e.g.
-
-```yaml
-ask_vqa_top:
-  bindings:
-    vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.top)"
-    vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)"
-  messages:
-    - role: user
-      stream: high_level
-      if_present: vqa_query
-      content:
-        - { type: image, feature: observation.images.top }
-        - { type: text, text: "${vqa_query}" }
-    - { role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa }
-```
-
-Add one such sub-recipe per camera the dataset records.
-
-## Recipe anatomy
-
-Recipes are YAML files backed by `TrainingRecipe` and `MessageTurn`.
-
-```yaml
-messages:
-  - { role: user, content: "${task}", stream: high_level }
-  - { role: assistant, content: "${subtask}", stream: low_level, target: true }
-```
-
-Rendered samples use HF-style chat messages plus LeRobot sidecars:
-
-```python
-sample["messages"]
-sample["message_streams"]
-sample["target_message_indices"]
-```
-
-The renderer does not apply a tokenizer chat template. Policy processors decide how to serialize the messages for their backbone.
-
-## Blends
-
-Blend recipes select one weighted sub-recipe deterministically from the sample index.
-The canonical `recipes/pi05_hirobot.yaml` combines memory updates, interjection responses, high-level subtask prediction, low-level execution, and VQA.
-
-## Graceful absence
-
-If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op.
-If an event-scoped branch is selected on a frame without the required event row, rendering returns `None`, allowing a loader to retry another sample.

docs/source/libero_plus.mdx

@@ -1,188 +0,0 @@
-# LIBERO-plus
-
-LIBERO-plus is a **robustness benchmark** for Vision-Language-Action (VLA) models built on top of [LIBERO](./libero). It systematically stress-tests policies by applying **seven independent perturbation dimensions** to the original LIBERO task set, exposing failure modes that standard benchmarks miss.
-
-- Paper: [In-depth Robustness Analysis of Vision-Language-Action Models](https://arxiv.org/abs/2510.13626)
-- GitHub: [sylvestf/LIBERO-plus](https://github.com/sylvestf/LIBERO-plus)
-- Dataset: [lerobot/libero_plus](https://huggingface.co/datasets/lerobot/libero_plus)
-
-![An overview of the LIBERO-plus benchmark perturbation dimensions](https://github.com/sylvestf/LIBERO-plus/raw/main/static/images/libero-plus.jpg)
-
-## Perturbation dimensions
-
-LIBERO-plus creates ~10 000 task variants by perturbing each original LIBERO task along these axes:
-
-| Dimension             | What changes                                          |
-| --------------------- | ----------------------------------------------------- |
-| Objects layout        | Target position, presence of confounding objects      |
-| Camera viewpoints     | Camera position, orientation, field-of-view           |
-| Robot initial states  | Manipulator start pose                                |
-| Language instructions | LLM-rewritten task description (paraphrase / synonym) |
-| Light conditions      | Intensity, direction, color, shadow                   |
-| Background textures   | Scene surface and object appearance                   |
-| Sensor noise          | Photometric distortions and image degradation         |
-
-## Available task suites
-
-LIBERO-plus covers the same five suites as LIBERO:
-
-| Suite          | CLI name         | Tasks | Max steps | Description                                        |
-| -------------- | ---------------- | ----- | --------- | -------------------------------------------------- |
-| LIBERO-Spatial | `libero_spatial` | 10    | 280       | Tasks requiring reasoning about spatial relations  |
-| LIBERO-Object  | `libero_object`  | 10    | 280       | Tasks centered on manipulating different objects   |
-| LIBERO-Goal    | `libero_goal`    | 10    | 300       | Goal-conditioned tasks with changing targets       |
-| LIBERO-90      | `libero_90`      | 90    | 400       | Short-horizon tasks from the LIBERO-100 collection |
-| LIBERO-Long    | `libero_10`      | 10    | 520       | Long-horizon tasks from the LIBERO-100 collection  |
-
-<Tip warning={true}>
-  Installing LIBERO-plus **replaces** vanilla LIBERO — it uninstalls `hf-libero`
-  so that `import libero` resolves to the LIBERO-plus fork. You cannot have both
-  installed at the same time. To switch back to vanilla LIBERO, uninstall the
-  fork and reinstall with `pip install -e ".[libero]"`.
-</Tip>
-
-## Installation
-
-### System dependencies (Linux only)
-
-```bash
-sudo apt install libexpat1 libfontconfig1-dev libmagickwand-dev
-```
-
-### Python package
-
-```bash
-pip install -e ".[libero]" "robosuite==1.4.1" bddl easydict mujoco wand scikit-image gym
-git clone https://github.com/sylvestf/LIBERO-plus.git
-cd LIBERO-plus && pip install --no-deps -e .
-pip uninstall -y hf-libero  # so `import libero` resolves to the fork
-```
-
-LIBERO-plus is installed from its GitHub fork rather than a pyproject extra — the fork ships as a namespace package that pip can't handle, so it must be cloned and added to `PYTHONPATH`. See `docker/Dockerfile.benchmark.libero_plus` for the canonical install. MuJoCo is required, so only Linux is supported.
-
-<Tip>
-Set the MuJoCo rendering backend before running evaluation:
-
-```bash
-export MUJOCO_GL=egl   # headless / HPC / cloud
-```
-
-</Tip>
-
-### Download LIBERO-plus assets
-
-LIBERO-plus ships its extended asset pack separately. Download `assets.zip` from the [Hugging Face dataset](https://huggingface.co/datasets/Sylvest/LIBERO-plus/tree/main) and extract it into the LIBERO-plus package directory:
-
-```bash
-# After installing the package, find where it was installed:
-python -c "import libero; print(libero.__file__)"
-# Then extract assets.zip into <package_root>/libero/assets/
-```
-
-## Evaluation
-
-### Default evaluation (recommended)
-
-Evaluate across the four standard suites (10 episodes per task):
-
-```bash
-lerobot-eval \
-  --policy.path="your-policy-id" \
-  --env.type=libero_plus \
-  --env.task=libero_spatial,libero_object,libero_goal,libero_10 \
-  --eval.batch_size=1 \
-  --eval.n_episodes=10 \
-  --env.max_parallel_tasks=1
-```
-
-### Single-suite evaluation
-
-Evaluate on one LIBERO-plus suite:
-
-```bash
-lerobot-eval \
-  --policy.path="your-policy-id" \
-  --env.type=libero_plus \
-  --env.task=libero_spatial \
-  --eval.batch_size=1 \
-  --eval.n_episodes=10
-```
-
-- `--env.task` picks the suite (`libero_spatial`, `libero_object`, etc.).
-- `--env.task_ids` restricts to specific task indices (`[0]`, `[1,2,3]`, etc.). Omit to run all tasks in the suite.
-- `--eval.batch_size` controls how many environments run in parallel.
-- `--eval.n_episodes` sets how many episodes to run per task.
-
-### Multi-suite evaluation
-
-Benchmark a policy across multiple suites at once by passing a comma-separated list:
-
-```bash
-lerobot-eval \
-  --policy.path="your-policy-id" \
-  --env.type=libero_plus \
-  --env.task=libero_spatial,libero_object \
-  --eval.batch_size=1 \
-  --eval.n_episodes=10
-```
-
-### Control mode
-
-LIBERO-plus supports two control modes — `relative` (default) and `absolute`. Different VLA checkpoints are trained with different action parameterizations, so make sure the mode matches your policy:
-
-```bash
---env.control_mode=relative   # or "absolute"
-```
-
-### Policy inputs and outputs
-
-**Observations:**
-
-- `observation.state` — 8-dim proprioceptive features (eef position, axis-angle orientation, gripper qpos)
-- `observation.images.image` — main camera view (`agentview_image`), HWC uint8
-- `observation.images.image2` — wrist camera view (`robot0_eye_in_hand_image`), HWC uint8
-
-**Actions:**
-
-- Continuous control in `Box(-1, 1, shape=(7,))` — 6D end-effector delta + 1D gripper
-
-### Recommended evaluation episodes
-
-For reproducible benchmarking, use **10 episodes per task** across all four standard suites (Spatial, Object, Goal, Long). This gives 400 total episodes and matches the protocol used for published results.
-
-## Training
-
-### Dataset
-
-A LeRobot-format training dataset for LIBERO-plus is available at:
-
-- [lerobot/libero_plus](https://huggingface.co/datasets/lerobot/libero_plus)
-
-### Example training command
-
-```bash
-lerobot-train \
-    --policy.type=smolvla \
-    --policy.repo_id=${HF_USER}/smolvla_libero_plus \
-    --policy.load_vlm_weights=true \
-    --dataset.repo_id=lerobot/libero_plus \
-    --env.type=libero_plus \
-    --env.task=libero_spatial \
-    --output_dir=./outputs/ \
-    --steps=100000 \
-    --batch_size=4 \
-    --eval.batch_size=1 \
-    --eval.n_episodes=1 \
-    --eval_freq=1000
-```
-
-## Relationship to LIBERO
-
-LIBERO-plus is a drop-in extension of LIBERO:
-
-- Same Python gym interface (`LiberoEnv`, `LiberoProcessorStep`)
-- Same camera names and observation/action format
-- Same task suite names
-- Installs under the same `libero` Python package name (different GitHub repo)
-
-To use the original LIBERO benchmark, see [LIBERO](./libero) and use `--env.type=libero`.

docs/source/robocasa.mdx

@@ -1,188 +0,0 @@
-# RoboCasa365
-
-[RoboCasa365](https://robocasa.ai) is a large-scale simulation framework for training and benchmarking **generalist robots** in everyday kitchen tasks. It ships 365 diverse manipulation tasks across 2,500 kitchen environments, 3,200+ object assets and 600+ hours of human demonstration data, on a PandaOmron 12-DOF mobile manipulator (Franka arm on a holonomic base).
-
-- Paper: [RoboCasa: Large-Scale Simulation of Everyday Tasks for Generalist Robots](https://arxiv.org/abs/2406.02523)
-- GitHub: [robocasa/robocasa](https://github.com/robocasa/robocasa)
-- Project website: [robocasa.ai](https://robocasa.ai)
-- Pretrained policy: [`lerobot/smolvla_robocasa`](https://huggingface.co/lerobot/smolvla_robocasa)
-- Single-task dataset (CloseFridge): [`pepijn223/robocasa_CloseFridge`](https://huggingface.co/datasets/pepijn223/robocasa_CloseFridge)
-
-<img
-  src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/robocasa-banner.webp"
-  alt="RoboCasa365 benchmark overview"
-  width="85%"
-/>
-
-## Available tasks
-
-RoboCasa365 organizes its 365 tasks into two families and three upstream benchmark groups that LeRobot exposes as first-class `--env.task` shortcuts:
-
-| Family    | Tasks | Description                                                                     |
-| --------- | ----- | ------------------------------------------------------------------------------- |
-| Atomic    | ~65   | Single-skill tasks: pick-and-place, door/drawer manipulation, appliance control |
-| Composite | ~300  | Multi-step tasks across 60+ categories: cooking, cleaning, organizing, etc.     |
-
-**Atomic task examples:** `CloseFridge`, `OpenDrawer`, `OpenCabinet`, `TurnOnMicrowave`, `TurnOffStove`, `NavigateKitchen`, `PickPlaceCounterToStove`.
-
-**Composite task categories:** baking, boiling, brewing, chopping, clearing table, defrosting food, loading dishwasher, making tea, microwaving food, washing dishes, and more.
-
-`--env.task` accepts three forms:
-
-- a single task name (`CloseFridge`)
-- a comma-separated list (`CloseFridge,OpenBlenderLid,PickPlaceCoffee`)
-- a benchmark-group shortcut — `atomic_seen`, `composite_seen`, `composite_unseen`, `pretrain50`, `pretrain100`, `pretrain200`, `pretrain300` — which auto-expands to the upstream task list and auto-sets the dataset `split` (`target` or `pretrain`).
-
-## Installation
-
-RoboCasa and its dependency `robosuite` are not published on PyPI, and RoboCasa's own `setup.py` hardcodes `lerobot==0.3.3`, which conflicts with this repo's `lerobot`. LeRobot therefore does **not** expose a `robocasa` extra — install the two packages manually as editable clones (using `--no-deps` on `robocasa` to skip its shadowed `lerobot` pin):
-
-```bash
-# After following the standard LeRobot installation instructions.
-
-git clone https://github.com/robocasa/robocasa.git ~/robocasa
-git clone https://github.com/ARISE-Initiative/robosuite.git ~/robosuite
-pip install -e ~/robocasa --no-deps
-pip install -e ~/robosuite
-
-# Robocasa's runtime deps (the ones its setup.py would have pulled, minus
-# the bad lerobot pin).
-pip install numpy numba scipy mujoco pygame Pillow opencv-python \
-            pyyaml pynput tqdm termcolor imageio h5py lxml hidapi \
-            tianshou gymnasium
-
-python -m robocasa.scripts.setup_macros
-# Lightweight assets (lightwheel object meshes + textures). Enough for
-# the default env out of the box.
-python -m robocasa.scripts.download_kitchen_assets \
-  --type tex tex_generative fixtures_lw objs_lw
-# Optional: full objaverse/aigen registries (~30GB) for richer object
-# variety. Enable at eval time via --env.obj_registries (see below).
-# python -m robocasa.scripts.download_kitchen_assets --type objs_objaverse
-```
-
-<Tip>
-RoboCasa requires MuJoCo. Set the rendering backend before training or evaluation:
-
-```bash
-export MUJOCO_GL=egl  # for headless servers (HPC, cloud)
-```
-
-</Tip>
-
-### Object registries
-
-By default the env samples objects only from the `lightwheel` registry (what `--type objs_lw` ships), which avoids a `Probabilities contain NaN` crash when the objaverse / aigen packs aren't on disk. If you've downloaded the full asset set, enable the full registry at runtime:
-
-```bash
---env.obj_registries='[objaverse,lightwheel]'
-```
-
-## Evaluation
-
-All eval snippets below mirror the CI command (see `.github/workflows/benchmark_tests.yml`). The `--rename_map` argument maps RoboCasa's native camera keys (`robot0_agentview_left` / `robot0_eye_in_hand` / `robot0_agentview_right`) onto the three-camera (`camera1` / `camera2` / `camera3`) input layout the released `smolvla_robocasa` policy was trained on.
-
-### Single-task evaluation (recommended for quick iteration)
-
-```bash
-lerobot-eval \
-  --policy.path=lerobot/smolvla_robocasa \
-  --env.type=robocasa \
-  --env.task=CloseFridge \
-  --eval.batch_size=1 \
-  --eval.n_episodes=20 \
-  --eval.use_async_envs=false \
-  --policy.device=cuda \
-  '--rename_map={"observation.images.robot0_agentview_left": "observation.images.camera1", "observation.images.robot0_eye_in_hand": "observation.images.camera2", "observation.images.robot0_agentview_right": "observation.images.camera3"}'
-```
-
-### Multi-task evaluation
-
-Pass a comma-separated list of tasks:
-
-```bash
-lerobot-eval \
-  --policy.path=lerobot/smolvla_robocasa \
-  --env.type=robocasa \
-  --env.task=CloseFridge,OpenCabinet,OpenDrawer,TurnOnMicrowave,TurnOffStove \
-  --eval.batch_size=1 \
-  --eval.n_episodes=20 \
-  --eval.use_async_envs=false \
-  --policy.device=cuda \
-  '--rename_map={"observation.images.robot0_agentview_left": "observation.images.camera1", "observation.images.robot0_eye_in_hand": "observation.images.camera2", "observation.images.robot0_agentview_right": "observation.images.camera3"}'
-```
-
-### Benchmark-group evaluation
-
-Run an entire upstream group (e.g. all 18 `atomic_seen` tasks with `split=target`):
-
-```bash
-lerobot-eval \
-  --policy.path=lerobot/smolvla_robocasa \
-  --env.type=robocasa \
-  --env.task=atomic_seen \
-  --eval.batch_size=1 \
-  --eval.n_episodes=20 \
-  --eval.use_async_envs=false \
-  --policy.device=cuda \
-  '--rename_map={"observation.images.robot0_agentview_left": "observation.images.camera1", "observation.images.robot0_eye_in_hand": "observation.images.camera2", "observation.images.robot0_agentview_right": "observation.images.camera3"}'
-```
-
-### Recommended evaluation episodes
-
-**20 episodes per task** for reproducible benchmarking. Matches the protocol used in published results.
-
-## Policy inputs and outputs
-
-**Observations** (raw RoboCasa camera names are preserved verbatim):
-
-- `observation.state` — 16-dim proprioceptive state (base position, base quaternion, relative end-effector position, relative end-effector quaternion, gripper qpos)
-- `observation.images.robot0_agentview_left` — left agent view, 256×256 HWC uint8
-- `observation.images.robot0_eye_in_hand` — wrist camera view, 256×256 HWC uint8
-- `observation.images.robot0_agentview_right` — right agent view, 256×256 HWC uint8
-
-**Actions:**
-
-- Continuous control in `Box(-1, 1, shape=(12,))` — base motion (4D) + control mode (1D) + end-effector position (3D) + end-effector rotation (3D) + gripper (1D).
-
-## Training
-
-### Single-task example
-
-A ready-to-use single-task dataset is on the Hub:
-[`pepijn223/robocasa_CloseFridge`](https://huggingface.co/datasets/pepijn223/robocasa_CloseFridge).
-
-Fine-tune a SmolVLA base on `CloseFridge`:
-
-```bash
-lerobot-train \
-  --policy.type=smolvla \
-  --policy.repo_id=${HF_USER}/smolvla_robocasa_CloseFridge \
-  --policy.load_vlm_weights=true \
-  --policy.push_to_hub=true \
-  --dataset.repo_id=pepijn223/robocasa_CloseFridge \
-  --env.type=robocasa \
-  --env.task=CloseFridge \
-  --output_dir=./outputs/smolvla_robocasa_CloseFridge \
-  --steps=100000 \
-  --batch_size=4 \
-  --eval_freq=5000 \
-  --eval.batch_size=1 \
-  --eval.n_episodes=5 \
-  --save_freq=10000
-```
-
-Evaluate the resulting checkpoint:
-
-```bash
-lerobot-eval \
-  --policy.path=${HF_USER}/smolvla_robocasa_CloseFridge \
-  --env.type=robocasa \
-  --env.task=CloseFridge \
-  --eval.batch_size=1 \
-  --eval.n_episodes=20
-```
-
-## Reproducing published results
-
-The released checkpoint [`lerobot/smolvla_robocasa`](https://huggingface.co/lerobot/smolvla_robocasa) is evaluated with the commands in the [Evaluation](#evaluation) section. CI runs a 10-atomic-task smoke eval (one episode each) on every PR touching the benchmark, picking fixture-centric tasks that don't require the objaverse asset pack.

docs/source/robocerebra.mdx

@@ -1,99 +0,0 @@
-# RoboCerebra
-
-[RoboCerebra](https://robocerebra-project.github.io/) is a long-horizon manipulation benchmark that evaluates **high-level reasoning, planning, and memory** in VLAs. Episodes chain multiple sub-goals with language-grounded intermediate instructions, built on top of LIBERO's simulator stack (MuJoCo + robosuite, Franka Panda 7-DOF).
-
-- Paper: [RoboCerebra: A Large-scale Benchmark for Long-horizon Robotic Manipulation Evaluation](https://arxiv.org/abs/2506.06677)
-- Project website: [robocerebra-project.github.io](https://robocerebra-project.github.io/)
-- Dataset: [`lerobot/robocerebra_unified`](https://huggingface.co/datasets/lerobot/robocerebra_unified) — LeRobot v3.0, 6,660 episodes / 571,116 frames at 20 fps, 1,728 language-grounded sub-tasks.
-- Pretrained policy: [`lerobot/smolvla_robocerebra`](https://huggingface.co/lerobot/smolvla_robocerebra)
-
-## Available tasks
-
-RoboCerebra reuses LIBERO's simulator, so evaluation runs against the LIBERO `libero_10` long-horizon suite:
-
-| Suite     | CLI name    | Tasks | Description                                                   |
-| --------- | ----------- | ----- | ------------------------------------------------------------- |
-| LIBERO-10 | `libero_10` | 10    | Long-horizon kitchen/living room tasks chaining 3–6 sub-goals |
-
-Each RoboCerebra episode in the dataset is segmented into multiple sub-tasks with natural-language instructions, which the unified dataset exposes as independent supervision signals.
-
-## Installation
-
-RoboCerebra piggybacks on LIBERO, so the `libero` extra is all you need:
-
-```bash
-pip install -e ".[libero]"
-```
-
-<Tip>
-RoboCerebra requires Linux (MuJoCo / robosuite). Set the rendering backend before training or evaluation:
-
-```bash
-export MUJOCO_GL=egl  # for headless servers (HPC, cloud)
-```
-
-</Tip>
-
-## Evaluation
-
-RoboCerebra eval runs against LIBERO's `libero_10` suite with RoboCerebra's camera naming (`image` + `wrist_image`) and an extra empty-camera slot so a three-view-trained policy receives the expected input layout:
-
-```bash
-lerobot-eval \
-  --policy.path=lerobot/smolvla_robocerebra \
-  --env.type=libero \
-  --env.task=libero_10 \
-  --env.fps=20 \
-  --env.obs_type=pixels_agent_pos \
-  --env.observation_height=256 \
-  --env.observation_width=256 \
-  '--env.camera_name_mapping={"agentview_image": "image", "robot0_eye_in_hand_image": "wrist_image"}' \
-  --eval.batch_size=1 \
-  --eval.n_episodes=10 \
-  --eval.use_async_envs=false \
-  --policy.device=cuda \
-  '--rename_map={"observation.images.image": "observation.images.camera1", "observation.images.wrist_image": "observation.images.camera2"}' \
-  --policy.empty_cameras=1
-```
-
-### Recommended evaluation episodes
-
-**10 episodes per task** across the `libero_10` suite (100 total) for reproducible benchmarking. Matches the protocol used in the RoboCerebra paper.
-
-## Policy inputs and outputs
-
-**Observations:**
-
-- `observation.state` — 8-dim proprioceptive state (7 joint positions + gripper)
-- `observation.images.image` — third-person view, 256×256 HWC uint8
-- `observation.images.wrist_image` — wrist-mounted camera view, 256×256 HWC uint8
-
-**Actions:**
-
-- Continuous control in `Box(-1, 1, shape=(7,))` — end-effector delta (6D) + gripper (1D)
-
-## Training
-
-The unified dataset at [`lerobot/robocerebra_unified`](https://huggingface.co/datasets/lerobot/robocerebra_unified) exposes two RGB streams and language-grounded sub-task annotations:
-
-| Feature                          | Shape         | Description          |
-| -------------------------------- | ------------- | -------------------- |
-| `observation.images.image`       | (256, 256, 3) | Third-person view    |
-| `observation.images.wrist_image` | (256, 256, 3) | Wrist-mounted camera |
-| `observation.state`              | (8,)          | Joint pos + gripper  |
-| `action`                         | (7,)          | EEF delta + gripper  |
-
-Fine-tune a SmolVLA base on it:
-
-```bash
-lerobot-train \
-  --policy.path=lerobot/smolvla_base \
-  --dataset.repo_id=lerobot/robocerebra_unified \
-  --env.type=libero \
-  --env.task=libero_10 \
-  --output_dir=outputs/smolvla_robocerebra
-```
-
-## Reproducing published results
-
-The released checkpoint [`lerobot/smolvla_robocerebra`](https://huggingface.co/lerobot/smolvla_robocerebra) was trained on `lerobot/robocerebra_unified` and evaluated with the command in the [Evaluation](#evaluation) section. CI runs the same command with `--eval.n_episodes=1` as a smoke test on every PR touching the benchmark.

docs/source/robomme.mdx

@@ -1,130 +0,0 @@
-# RoboMME
-
-[RoboMME](https://robomme.github.io) is a memory-augmented manipulation benchmark built on ManiSkill (SAPIEN). It evaluates a robot's ability to retain and use information across an episode — counting, object permanence, reference, and imitation.
-
-- **16 tasks** across 4 memory-skill suites
-- **1,600 training demos** (100 per task, 50 val, 50 test)
-- **Dataset**: [`lerobot/robomme`](https://huggingface.co/datasets/lerobot/robomme) — LeRobot v3.0, 768K frames at 10 fps
-- **Simulator**: ManiSkill / SAPIEN, Panda arm, Linux only
-
-![RoboMME benchmark tasks overview](https://cdn-thumbnails.huggingface.co/social-thumbnails/papers/2603.04639/gradient.png)
-
-## Tasks
-
-| Suite                             | Tasks                                                         |
-| --------------------------------- | ------------------------------------------------------------- |
-| **Counting** (temporal memory)    | BinFill, PickXtimes, SwingXtimes, StopCube                    |
-| **Permanence** (spatial memory)   | VideoUnmask, VideoUnmaskSwap, ButtonUnmask, ButtonUnmaskSwap  |
-| **Reference** (object memory)     | PickHighlight, VideoRepick, VideoPlaceButton, VideoPlaceOrder |
-| **Imitation** (procedural memory) | MoveCube, InsertPeg, PatternLock, RouteStick                  |
-
-## Installation
-
-> RoboMME requires **Linux** (ManiSkill/SAPIEN uses Vulkan rendering). Docker is recommended to isolate dependency conflicts.
-
-### Native (Linux)
-
-```bash
-pip install --override <(printf 'gymnasium==0.29.1\nnumpy==1.26.4\n') \
-  -e '.[smolvla,av-dep]' \
-  'robomme @ git+https://github.com/RoboMME/robomme_benchmark.git@main'
-```
-
-> **Dependency note**: `mani-skill` (pulled by `robomme`) pins `gymnasium==0.29.1` and `numpy<2.0.0`, which conflict with lerobot's base `numpy>=2.0.0`. That's why `robomme` is not a pyproject extra — use the override install above, or the Docker approach below to avoid conflicts entirely.
-
-### Docker (recommended)
-
-```bash
-# Build base image first (from repo root)
-docker build -f docker/Dockerfile.eval-base -t lerobot-eval-base .
-
-# Build RoboMME eval image (applies gymnasium + numpy pin overrides)
-docker build -f docker/Dockerfile.benchmark.robomme -t lerobot-robomme .
-```
-
-The `docker/Dockerfile.benchmark.robomme` image overrides `gymnasium==0.29.1` and `numpy==1.26.4` after lerobot's install. Both versions are runtime-safe for lerobot's actual API usage.
-
-## Running Evaluation
-
-### Default (single task, single episode)
-
-```bash
-lerobot-eval \
-    --policy.path=<your_policy_repo> \
-    --env.type=robomme \
-    --env.task=PickXtimes \
-    --env.dataset_split=test \
-    --env.task_ids=[0] \
-    --eval.batch_size=1 \
-    --eval.n_episodes=1
-```
-
-### Multi-task evaluation
-
-Evaluate multiple tasks in one run by comma-separating task names. Use `task_ids` to control which episodes are evaluated per task. Recommended: 50 episodes per task for the test split.
-
-```bash
-lerobot-eval \
-    --policy.path=<your_policy_repo> \
-    --env.type=robomme \
-    --env.task=PickXtimes,BinFill,StopCube,MoveCube,InsertPeg \
-    --env.dataset_split=test \
-    --env.task_ids=[0,1,2,3,4,5,6,7,8,9] \
-    --eval.batch_size=1 \
-    --eval.n_episodes=50
-```
-
-### Key CLI options for `env.type=robomme`
-
-| Option               | Default       | Description                                        |
-| -------------------- | ------------- | -------------------------------------------------- |
-| `env.task`           | `PickXtimes`  | Any of the 16 task names above (comma-separated)   |
-| `env.dataset_split`  | `test`        | `train`, `val`, or `test`                          |
-| `env.action_space`   | `joint_angle` | `joint_angle` (8-D) or `ee_pose` (7-D)             |
-| `env.episode_length` | `300`         | Max steps per episode                              |
-| `env.task_ids`       | `null`        | List of episode indices to evaluate (null = `[0]`) |
-
-## Dataset
-
-The dataset [`lerobot/robomme`](https://huggingface.co/datasets/lerobot/robomme) is in **LeRobot v3.0 format** and can be loaded directly:
-
-```python
-from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-dataset = LeRobotDataset("lerobot/robomme")
-```
-
-### Dataset features
-
-| Feature            | Shape         | Description                     |
-| ------------------ | ------------- | ------------------------------- |
-| `image`            | (256, 256, 3) | Front camera RGB                |
-| `wrist_image`      | (256, 256, 3) | Wrist camera RGB                |
-| `actions`          | (8,)          | Joint angles + gripper          |
-| `state`            | (8,)          | Joint positions + gripper state |
-| `simple_subgoal`   | str           | High-level language annotation  |
-| `grounded_subgoal` | str           | Grounded language annotation    |
-| `episode_index`    | int           | Episode ID                      |
-| `frame_index`      | int           | Frame within episode            |
-
-### Feature key alignment (training)
-
-The env wrapper exposes `pixels/image` and `pixels/wrist_image` as observation keys. The `features_map` in `RoboMMEEnv` maps these to `observation.images.image` and `observation.images.wrist_image` for the policy. State is exposed as `agent_pos` and maps to `observation.state`.
-
-The dataset's `image` and `wrist_image` columns already align with the policy input keys, so no renaming is needed when fine-tuning.
-
-## Action Spaces
-
-| Type          | Dim | Description                                               |
-| ------------- | --- | --------------------------------------------------------- |
-| `joint_angle` | 8   | 7 joint angles + 1 gripper (−1 closed, +1 open, absolute) |
-| `ee_pose`     | 7   | xyz + roll/pitch/yaw + gripper                            |
-
-Set via `--env.action_space=joint_angle` (default) or `--env.action_space=ee_pose`.
-
-## Platform Notes
-
-- **Linux only**: ManiSkill requires SAPIEN/Vulkan. macOS and Windows are not supported.
-- **GPU recommended**: Rendering is CPU-capable but slow; CUDA + Vulkan gives full speed.
-- **gymnasium / numpy conflict**: See installation note above. Docker image handles this automatically.
-- **ManiSkill fork**: `robomme` depends on a specific ManiSkill fork (`YinpeiDai/ManiSkill`), pulled in automatically via the `robomme` package.

docs/source/robotwin.mdx

@@ -1,223 +0,0 @@
-# RoboTwin 2.0
-
-RoboTwin 2.0 is a **large-scale dual-arm manipulation benchmark** built on the SAPIEN physics engine. It provides a standardized evaluation protocol for bimanual robotic policies across 50 tasks (as of upstream `main`) with strong domain randomization (clutter, lighting, background, tabletop height, and language instructions).
-
-- Paper: [RoboTwin 2.0: A Scalable Data Generator and Benchmark with Strong Domain Randomization for Robust Bimanual Robotic Manipulation](https://arxiv.org/abs/2506.18088)
-- GitHub: [RoboTwin-Platform/RoboTwin](https://github.com/RoboTwin-Platform/RoboTwin)
-- Leaderboard: [robotwin-platform.github.io/leaderboard](https://robotwin-platform.github.io/leaderboard)
-- Dataset: [lerobot/robotwin_unified](https://huggingface.co/datasets/lerobot/robotwin_unified)
-
-![RoboTwin 2.0 benchmark overview](https://www.aitntnews.com/pictures/2025/7/8/9a7f79cb-5ba9-11f0-8581-fa163e47d677.png)
-
-## Overview
-
-| Property      | Value                                                    |
-| ------------- | -------------------------------------------------------- |
-| Tasks         | 50 dual-arm manipulation tasks                           |
-| Robot         | Aloha-AgileX bimanual (14 DOF, 7 per arm)                |
-| Action space  | 14-dim joint-space, continuous in `[-1, 1]`              |
-| Cameras       | `head_camera`, `left_camera`, `right_camera`             |
-| Simulator     | SAPIEN (not MuJoCo)                                      |
-| Eval protocol | 100 episodes/task, 50 demo_clean demonstrations          |
-| Eval settings | **Easy** (`demo_clean`) and **Hard** (`demo_randomized`) |
-
-## Available tasks
-
-RoboTwin 2.0 ships 50 dual-arm manipulation tasks in its upstream `envs/` directory. The canonical list is the `ROBOTWIN_TASKS` tuple in `src/lerobot/envs/robotwin.py`, mirrored verbatim from the upstream repo. Example tasks:
-
-| Task                     | CLI name                 | Category          |
-| ------------------------ | ------------------------ | ----------------- |
-| Beat block with hammer   | `beat_block_hammer`      | Tool use          |
-| Click bell / alarm clock | `click_bell`             | Precision press   |
-| Stack blocks (2 / 3)     | `stack_blocks_two/three` | Stacking          |
-| Stack bowls (2 / 3)      | `stack_bowls_two/three`  | Stacking          |
-| Handover block / mic     | `handover_block`         | Bimanual coord.   |
-| Lift pot                 | `lift_pot`               | Bimanual lift     |
-| Shake bottle             | `shake_bottle`           | Continuous motion |
-| Turn switch              | `turn_switch`            | Articulated obj   |
-| Stamp seal               | `stamp_seal`             | Precision place   |
-| Scan object              | `scan_object`            | Mobile manip.     |
-
-Pass a comma-separated list to `--env.task` to run multiple tasks in a single eval sweep.
-
-<Tip warning={true}>
-  `open_laptop` is currently broken upstream (its `check_success()` uses
-  `self.arm_tag`, which is only set inside the scripted-expert `play_once()`
-  path and therefore unavailable during normal policy eval). Avoid it until the
-  upstream bug is fixed, or patch the task to default `self.arm_tag = "left"` in
-  `load_actors()`.
-</Tip>
-
-## Dataset
-
-The RoboTwin 2.0 dataset is available in **LeRobot v3.0 format** on the Hugging Face Hub:
-
-```
-lerobot/robotwin_unified
-```
-
-It contains over 100,000 pre-collected trajectories across all 50 tasks (79.6 GB, Apache 2.0 license). No format conversion is needed — it is already in the correct LeRobot v3.0 schema with video observations and action labels.
-
-You can load it directly with the HF Datasets library:
-
-```python
-from datasets import load_dataset
-
-ds = load_dataset("lerobot/robotwin_unified", split="train")
-```
-
-## Installation
-
-RoboTwin 2.0 requires **Linux** with an NVIDIA GPU (CUDA 12.1 recommended). Installation takes approximately 20 minutes.
-
-### 1. Create a conda environment
-
-```bash
-conda create -n robotwin python=3.10 -y
-conda activate robotwin
-```
-
-### 2. Install LeRobot
-
-```bash
-git clone https://github.com/huggingface/lerobot.git
-cd lerobot
-pip install -e "."
-```
-
-### 3. Install RoboTwin 2.0
-
-```bash
-git clone https://github.com/RoboTwin-Platform/RoboTwin.git
-cd RoboTwin
-bash script/_install.sh
-bash script/_download_assets.sh
-```
-
-The install script handles all Python dependencies including SAPIEN, CuRobo, mplib, and pytorch3d.
-
-<Tip warning={true}>
-If the automated install fails, install manually:
-
-```bash
-pip install -r requirements.txt
-pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
-cd envs && git clone https://github.com/NVlabs/curobo.git && cd curobo
-pip install -e . --no-build-isolation
-```
-
-Then apply the required mplib fix: in `mplib/planner.py` line 807, remove `or collide` from the conditional.
-
-</Tip>
-
-### 4. Add RoboTwin to PYTHONPATH
-
-The RoboTwin task modules must be importable by LeRobot. From within the `RoboTwin/` directory:
-
-```bash
-export PYTHONPATH="${PYTHONPATH}:$(pwd)"
-```
-
-Add this to your shell profile to make it permanent.
-
-## Evaluation
-
-### Standard evaluation (recommended)
-
-Evaluate a policy on a single task with the official protocol (100 episodes):
-
-```bash
-lerobot-eval \
-  --policy.path="your-hf-policy-id" \
-  --env.type=robotwin \
-  --env.task=beat_block_hammer \
-  --eval.batch_size=1 \
-  --eval.n_episodes=100
-```
-
-### Single-task quick check
-
-```bash
-lerobot-eval \
-  --policy.path="your-hf-policy-id" \
-  --env.type=robotwin \
-  --env.task=beat_block_hammer \
-  --eval.batch_size=1 \
-  --eval.n_episodes=5
-```
-
-### Multi-task sweep
-
-Evaluate on several tasks in one run:
-
-```bash
-lerobot-eval \
-  --policy.path="your-hf-policy-id" \
-  --env.type=robotwin \
-  --env.task=beat_block_hammer,click_bell,handover_block,stack_blocks_two \
-  --eval.batch_size=1 \
-  --eval.n_episodes=100
-```
-
-### Full benchmark (all 50 tasks)
-
-```bash
-lerobot-eval \
-  --policy.path="your-hf-policy-id" \
-  --env.type=robotwin \
-  --env.task=adjust_bottle,beat_block_hammer,blocks_ranking_rgb,blocks_ranking_size,click_alarmclock,click_bell,dump_bin_bigbin,grab_roller,handover_block,handover_mic,hanging_mug,lift_pot,move_can_pot,move_pillbottle_pad,move_playingcard_away,move_stapler_pad,open_microwave,pick_diverse_bottles,pick_dual_bottles,place_a2b_left,place_a2b_right,place_bread_basket,place_bread_skillet,place_burger_fries,place_can_basket,place_cans_plasticbox,place_container_plate,place_dual_shoes,place_empty_cup,place_fan,place_mouse_pad,place_object_basket,place_object_scale,place_object_stand,place_phone_stand,place_shoe,press_stapler,put_bottles_dustbin,put_object_cabinet,rotate_qrcode,scan_object,shake_bottle,shake_bottle_horizontally,stack_blocks_three,stack_blocks_two,stack_bowls_three,stack_bowls_two,stamp_seal,turn_switch \
-  --eval.batch_size=1 \
-  --eval.n_episodes=100
-```
-
-<Tip>
-  `open_laptop` is intentionally omitted above because of the upstream
-  `self.arm_tag` bug (see the **Available tasks** section). Re-add it once the
-  upstream fix lands.
-</Tip>
-
-## Camera configuration
-
-By default, all three cameras are included:
-
-| Camera key     | Description                    |
-| -------------- | ------------------------------ |
-| `head_camera`  | Torso-mounted overhead view    |
-| `left_camera`  | Left arm wrist-mounted camera  |
-| `right_camera` | Right arm wrist-mounted camera |
-
-To use a subset of cameras, override `--env.camera_names`:
-
-```bash
-lerobot-eval \
-  --policy.path="your-hf-policy-id" \
-  --env.type=robotwin \
-  --env.task=beat_block_hammer \
-  --env.camera_names="head_camera,left_camera" \
-  --eval.batch_size=1 \
-  --eval.n_episodes=10
-```
-
-## Environment config reference
-
-Key parameters for `RoboTwinEnvConfig`:
-
-| Parameter            | Default                                  | Description                        |
-| -------------------- | ---------------------------------------- | ---------------------------------- |
-| `task`               | `"beat_block_hammer"`                    | Comma-separated task name(s)       |
-| `fps`                | `25`                                     | Simulation FPS                     |
-| `episode_length`     | `300`                                    | Max steps per episode              |
-| `obs_type`           | `"pixels_agent_pos"`                     | `"pixels"` or `"pixels_agent_pos"` |
-| `camera_names`       | `"head_camera,left_camera,right_camera"` | Comma-separated active cameras     |
-| `observation_height` | `240`                                    | Camera pixel height                |
-| `observation_width`  | `320`                                    | Camera pixel width                 |
-
-## Leaderboard submission
-
-Results can be submitted to the [RoboTwin 2.0 leaderboard](https://robotwin-platform.github.io/leaderboard). The official protocol requires:
-
-- Training on 50 `demo_clean` demonstrations per task
-- Evaluating 100 episodes per task
-- Reporting success rate separately for **Easy** (`demo_clean`) and **Hard** (`demo_randomized`) settings
-
-For submission instructions, refer to the [RoboTwin 2.0 documentation](https://robotwin-platform.github.io/doc/).

docs/source/rtc.mdx

@@ -34,7 +34,7 @@ pip install -e ".[smolvla]"
 
 ### Using RTC with Pi0
 
-You can find a complete reference implementation in [eval_with_real_robot.py](examples/rtc/eval_with_real_robot.py).
+You can use `lerobot-rollout --strategy.type=base --inference.type=rtc` for RTC deployment on real robots.
 The snippet below provides a simplified pseudo-example of how RTC operates with Pi0 in your pipeline:
 
 ```python
@@ -137,8 +137,12 @@ The script generates a visualization of the denoising process, comparing standar
 ## Testing RTC with a Real Robot
 
 ```bash
-python examples/rtc/eval_with_real_robot.py \
+lerobot-rollout \
+    --strategy.type=base \
     --policy.path=${HF_USERNAME}/policy_repo_id \
+    --inference.type=rtc \
+    --inference.rtc.execution_horizon=10 \
+    --inference.rtc.max_guidance_weight=10.0 \
     --robot.type=so100_follower \
     --robot.port=/dev/tty.usbmodem58FA0834591 \
     --robot.cameras="{ gripper: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
@@ -178,7 +182,7 @@ visualizer = RTCDebugVisualizer()
 # ... create plots
 ```
 
-See `examples/rtc/eval_dataset.py` for a complete example of visualization.
+See `examples/rtc/eval_dataset.py` for a complete example of offline RTC visualization.
 
 ## References
 

docs/source/tools.mdx

@@ -1,198 +0,0 @@
-# Tools
-
-LeRobot v3.1 supports **tool calls** in policies — assistant messages can
-emit structured invocations like `say(text="OK, starting now")` that the
-runtime dispatches to a real implementation (TTS, controller, logger, …).
-
-This page covers:
-
-1. Where the tool catalog lives (PR 1).
-2. How the annotation pipeline produces tool-call atoms (PR 2).
-3. How to add your own tool (PR 3).
-
-## Where tools are declared
-
-Two layers.
-
-**The catalog** — a list of OpenAI-style function schemas — lives at
-`meta/info.json["tools"]` on each dataset. Example:
-
-```json
-{
-  "features": { "...": "..." },
-  "tools": [
-    {
-      "type": "function",
-      "function": {
-        "name": "say",
-        "description": "Speak a short utterance to the user via the TTS executor.",
-        "parameters": {
-          "type": "object",
-          "properties": {
-            "text": { "type": "string", "description": "The verbatim text to speak." }
-          },
-          "required": ["text"]
-        }
-      }
-    }
-  ]
-}
-```
-
-Read it via the dataset metadata accessor:
-
-```python
-from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
-
-meta = LeRobotDatasetMetadata(repo_id="pepijn/super_poulain_final_annotations")
-tools = meta.tools     # list[dict] — OpenAI tool schemas
-```
-
-If the dataset's `info.json` doesn't declare any tools, `meta.tools`
-returns `DEFAULT_TOOLS` from `lerobot.datasets.language` — currently a
-single-entry list with the canonical `say` schema. So unannotated
-datasets and chat-template consumers keep working without any
-configuration:
-
-```python
-prompt_str = tokenizer.apply_chat_template(
-    sample["messages"],
-    tools=meta.tools,                 # works either way
-    add_generation_prompt=False,
-    tokenize=False,
-)
-```
-
-**The implementations** — runnable Python — live under
-`src/lerobot/tools/`, one file per tool. The `say` implementation
-arrives in PR 3 and wraps Kyutai's pocket-tts model.
-
-## Per-row tool *invocations*
-
-The catalog above describes *what can be called*. The actual *call* — the
-function name plus the argument values — is stored per-row, on the
-assistant atoms in `language_events`:
-
-```python
-{
-  "role": "assistant",
-  "content": null,
-  "style": null,
-  "timestamp": 12.4,
-  "camera": null,
-  "tool_calls": [
-    { "type": "function",
-      "function": { "name": "say", "arguments": { "text": "On it." } } }
-  ]
-}
-```
-
-Recipes splice these into rendered messages via `tool_calls_from`:
-
-```yaml
-user_interjection_response:
-  bindings:
-    speech: "emitted_at(t, role=assistant, tool_name=say)"
-  messages:
-    - { role: user,      content: "${task}",         stream: high_level }
-    - { role: assistant, content: "${current_plan}", stream: high_level,
-        target: true, tool_calls_from: speech }
-```
-
-The model's training target is one assistant turn that carries both the
-plan text *and* the `say` tool call. At inference, the runtime parses
-the generated text back into structured `tool_calls` and dispatches to
-the matching implementation.
-
-## How to add your own tool
-
-Three steps. Concrete example: a `record_observation` tool the policy
-can call to capture an extra observation outside the regular control
-loop.
-
-### Step 1 — declare the schema
-
-Add an entry under `meta/info.json["tools"]`. Either edit the file
-directly on disk *before* running the annotation pipeline (it'll be
-preserved) or hand it to `lerobot-annotate` via a config flag (PR 2 —
-exact CLI lands with the pipeline change).
-
-```json
-{
-  "tools": [
-    { "type": "function", "function": { "name": "say", "...": "..." } },
-    {
-      "type": "function",
-      "function": {
-        "name": "record_observation",
-        "description": "Capture a high-resolution still image for the user.",
-        "parameters": {
-          "type": "object",
-          "properties": {
-            "label": { "type": "string", "description": "Short label for the saved image." }
-          },
-          "required": ["label"]
-        }
-      }
-    }
-  ]
-}
-```
-
-The schema follows OpenAI's function-calling convention exactly, so the
-chat template can render it natively.
-
-### Step 2 — implement the call
-
-Create `src/lerobot/tools/record_observation.py`:
-
-```python
-from .base import Tool
-from typing import Any
-
-RECORD_OBSERVATION_SCHEMA: dict[str, Any] = { "...": "..." }   # mirrors the JSON above
-
-
-class RecordObservationTool:
-    name = "record_observation"
-    schema = RECORD_OBSERVATION_SCHEMA
-
-    def __init__(self, schema: dict | None = None, output_dir: str = "."):
-        self.output_dir = output_dir
-
-    def call(self, arguments: dict) -> str:
-        label = arguments["label"]
-        # ... save the latest camera frame to <output_dir>/<label>.png ...
-        return f"saved {label}.png"
-```
-
-One file per tool keeps dependencies isolated — `record_observation`
-might pull `pillow`, while `say` (PR 3) pulls `pocket-tts`. Users
-installing only the tools they need avoid heavy transitive deps.
-
-### Step 3 — register it
-
-Add to `src/lerobot/tools/registry.py` (PR 3):
-
-```python
-from .record_observation import RecordObservationTool
-
-TOOL_REGISTRY["record_observation"] = RecordObservationTool
-```
-
-That's it. At runtime `get_tools(meta)` looks up each schema in
-`meta.tools`, instantiates the matching registered class, and returns
-a name → instance dict the dispatcher can route into.
-
-## Where this fits in the three-PR stack
-
-| Layer | PR | What lands |
-|---|---|---|
-| Catalog storage in `meta/info.json` + `meta.tools` accessor | PR 1 | This page; `SAY_TOOL_SCHEMA`, `DEFAULT_TOOLS` constants in `lerobot.datasets.language`; `LeRobotDatasetMetadata.tools` property |
-| Annotation pipeline writes `tools` to meta after a run; honors anything users pre-populated | PR 2 | `lerobot-annotate` ensures `meta/info.json["tools"]` includes the canonical `say` and merges any user-declared tools |
-| Runnable implementations under `src/lerobot/tools/`; runtime dispatcher; `say.py` wired to Kyutai's pocket-tts | PR 3 | One file per tool; `Tool` protocol; `TOOL_REGISTRY`; optional `[tools]` extra in `pyproject.toml` |
-
-If you want to use a tool *without* writing an implementation (e.g. for
-training-time chat-template formatting only), step 1 alone is enough —
-the model still learns to *generate* the call. Steps 2 and 3 are only
-needed to actually *execute* it at inference.

docs/source/unitree_g1.mdx

@@ -284,7 +284,7 @@ python examples/rtc/eval_with_real_robot.py \
   --task="task_description" \
   --duration=1000 \
   --fps=30 \
-  --rtc.enabled=true
+  --inference.type=rtc
 ```
 
 ---

docs/source/vlabench.mdx

@@ -1,176 +0,0 @@
-# VLABench
-
-[VLABench](https://github.com/OpenMOSS/VLABench) is a large-scale benchmark for **language-conditioned robotic manipulation with long-horizon reasoning**. The upstream suite covers 100 task categories across 2,000+ objects and evaluates six dimensions of robot intelligence: mesh & texture understanding, spatial reasoning, world-knowledge transfer, semantic instruction comprehension, physical-law understanding, and long-horizon planning. Built on MuJoCo / dm_control with a Franka Panda 7-DOF arm. LeRobot exposes **43 of these tasks** through `--env.task` (21 primitives + 22 composites, see [Available tasks](#available-tasks) below).
-
-- Paper: [VLABench: A Large-Scale Benchmark for Language-Conditioned Robotics Manipulation with Long-Horizon Reasoning](https://arxiv.org/abs/2412.18194)
-- GitHub: [OpenMOSS/VLABench](https://github.com/OpenMOSS/VLABench)
-- Project website: [vlabench.github.io](https://vlabench.github.io)
-- Pretrained policy: [`lerobot/smolvla_vlabench`](https://huggingface.co/lerobot/smolvla_vlabench)
-
-<img
-  src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/vlabench.png"
-  alt="VLABench benchmark overview"
-  width="85%"
-/>
-
-## Available tasks
-
-VLABench ships two task suites covering **43 task categories** in LeRobot's `--env.task` surface:
-
-| Suite     | CLI name    | Tasks | Description                                                      |
-| --------- | ----------- | ----- | ---------------------------------------------------------------- |
-| Primitive | `primitive` | 21    | Single / few-skill combinations (select, insert, physics QA)     |
-| Composite | `composite` | 22    | Multi-step reasoning and long-horizon planning (cook, rearrange) |
-
-**Primitive tasks:** `select_fruit`, `select_toy`, `select_chemistry_tube`, `add_condiment`, `select_book`, `select_painting`, `select_drink`, `insert_flower`, `select_billiards`, `select_ingredient`, `select_mahjong`, `select_poker`, and physical-reasoning tasks (`density_qa`, `friction_qa`, `magnetism_qa`, `reflection_qa`, `simple_cuestick_usage`, `simple_seesaw_usage`, `sound_speed_qa`, `thermal_expansion_qa`, `weight_qa`).
-
-**Composite tasks:** `cluster_billiards`, `cluster_book`, `cluster_drink`, `cluster_toy`, `cook_dishes`, `cool_drink`, `find_unseen_object`, `get_coffee`, `hammer_nail`, `heat_food`, `make_juice`, `play_mahjong`, `play_math_game`, `play_poker`, `play_snooker`, `rearrange_book`, `rearrange_chemistry_tube`, `set_dining_table`, `set_study_table`, `store_food`, `take_chemistry_experiment`, `use_seesaw_complex`.
-
-`--env.task` accepts three forms:
-
-- a single task name (`select_fruit`)
-- a comma-separated list (`select_fruit,heat_food`)
-- a suite shortcut (`primitive`, `composite`, or `primitive,composite`)
-
-## Installation
-
-VLABench is **not on PyPI** — its only distribution is the [OpenMOSS/VLABench](https://github.com/OpenMOSS/VLABench) GitHub repo — so LeRobot does not expose a `vlabench` extra. Install it manually as an editable clone, alongside the MuJoCo / dm_control pins VLABench needs, then fetch the mesh assets:
-
-```bash
-# After following the standard LeRobot installation instructions.
-
-git clone https://github.com/OpenMOSS/VLABench.git ~/VLABench
-git clone https://github.com/motion-planning/rrt-algorithms.git ~/rrt-algorithms
-pip install -e ~/VLABench -e ~/rrt-algorithms
-pip install "mujoco==3.2.2" "dm-control==1.0.22" \
-            open3d colorlog scikit-learn openai gdown
-
-python ~/VLABench/scripts/download_assets.py
-```
-
-<Tip>
-VLABench requires Linux (`sys_platform == 'linux'`) and Python 3.10+. Set the MuJoCo rendering backend before running:
-
-```bash
-export MUJOCO_GL=egl  # for headless servers (HPC, cloud)
-```
-
-</Tip>
-
-## Evaluation
-
-All eval snippets below mirror the command CI runs (see `.github/workflows/benchmark_tests.yml`). The `--rename_map` argument maps VLABench's `image` / `second_image` / `wrist_image` camera keys onto the three-camera (`camera1` / `camera2` / `camera3`) input layout the released `smolvla_vlabench` policy was trained on.
-
-### Single-task evaluation (recommended for quick iteration)
-
-```bash
-lerobot-eval \
-  --policy.path=lerobot/smolvla_vlabench \
-  --env.type=vlabench \
-  --env.task=select_fruit \
-  --eval.batch_size=1 \
-  --eval.n_episodes=10 \
-  --eval.use_async_envs=false \
-  --policy.device=cuda \
-  '--rename_map={"observation.images.image": "observation.images.camera1", "observation.images.second_image": "observation.images.camera2", "observation.images.wrist_image": "observation.images.camera3"}'
-```
-
-### Multi-task evaluation
-
-Pass a comma-separated list of tasks:
-
-```bash
-lerobot-eval \
-  --policy.path=lerobot/smolvla_vlabench \
-  --env.type=vlabench \
-  --env.task=select_fruit,select_toy,add_condiment,heat_food \
-  --eval.batch_size=1 \
-  --eval.n_episodes=10 \
-  --eval.use_async_envs=false \
-  --policy.device=cuda \
-  '--rename_map={"observation.images.image": "observation.images.camera1", "observation.images.second_image": "observation.images.camera2", "observation.images.wrist_image": "observation.images.camera3"}'
-```
-
-### Suite-wide evaluation
-
-Run an entire suite (all 21 primitives or all 22 composites):
-
-```bash
-lerobot-eval \
-  --policy.path=lerobot/smolvla_vlabench \
-  --env.type=vlabench \
-  --env.task=primitive \
-  --eval.batch_size=1 \
-  --eval.n_episodes=10 \
-  --eval.use_async_envs=false \
-  --policy.device=cuda \
-  --env.max_parallel_tasks=1 \
-  '--rename_map={"observation.images.image": "observation.images.camera1", "observation.images.second_image": "observation.images.camera2", "observation.images.wrist_image": "observation.images.camera3"}'
-```
-
-Or both suites:
-
-```bash
-lerobot-eval \
-  --policy.path=lerobot/smolvla_vlabench \
-  --env.type=vlabench \
-  --env.task=primitive,composite \
-  --eval.batch_size=1 \
-  --eval.n_episodes=10 \
-  --eval.use_async_envs=false \
-  --policy.device=cuda \
-  --env.max_parallel_tasks=1 \
-  '--rename_map={"observation.images.image": "observation.images.camera1", "observation.images.second_image": "observation.images.camera2", "observation.images.wrist_image": "observation.images.camera3"}'
-```
-
-### Recommended evaluation episodes
-
-**10 episodes per task** for reproducible benchmarking (210 total for the full primitive suite, 220 for composite). Matches the protocol in the VLABench paper.
-
-## Policy inputs and outputs
-
-**Observations:**
-
-- `observation.state` — 7-dim end-effector state (position xyz + Euler xyz + gripper)
-- `observation.images.image` — front camera, 480×480 HWC uint8
-- `observation.images.second_image` — second camera, 480×480 HWC uint8
-- `observation.images.wrist_image` — wrist camera, 480×480 HWC uint8
-
-**Actions:**
-
-- Continuous control in `Box(-1, 1, shape=(7,))` — 3D position + 3D Euler orientation + 1D gripper.
-
-## Training
-
-### Datasets
-
-Pre-collected VLABench datasets in LeRobot format on the Hub:
-
-- [`VLABench/vlabench_primitive_ft_lerobot_video`](https://huggingface.co/datasets/VLABench/vlabench_primitive_ft_lerobot_video) — 5,000 episodes, 128 tasks, 480×480 images.
-- [`VLABench/vlabench_composite_ft_lerobot_video`](https://huggingface.co/datasets/VLABench/vlabench_composite_ft_lerobot_video) — 5,977 episodes, 167 tasks, 224×224 images.
-
-### Example training command
-
-Fine-tune a SmolVLA base on the primitive suite:
-
-```bash
-lerobot-train \
-  --policy.type=smolvla \
-  --policy.repo_id=${HF_USER}/smolvla_vlabench_primitive \
-  --policy.load_vlm_weights=true \
-  --policy.push_to_hub=true \
-  --dataset.repo_id=VLABench/vlabench_primitive_ft_lerobot_video \
-  --env.type=vlabench \
-  --env.task=select_fruit \
-  --output_dir=./outputs/smolvla_vlabench_primitive \
-  --steps=100000 \
-  --batch_size=4 \
-  --eval_freq=5000 \
-  --eval.batch_size=1 \
-  --eval.n_episodes=1 \
-  --save_freq=10000
-```
-
-## Reproducing published results
-
-The released checkpoint [`lerobot/smolvla_vlabench`](https://huggingface.co/lerobot/smolvla_vlabench) was trained on the primitive-suite dataset above and is evaluated with the [Single-task](#single-task-evaluation-recommended-for-quick-iteration) / [Suite-wide](#suite-wide-evaluation) commands. CI runs a 10-primitive-task smoke eval (one episode each) on every PR touching the benchmark.

examples/annotation/run_hf_job.py

@@ -1,66 +0,0 @@
-#!/usr/bin/env python
-"""Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6 MoE).
-
-Spawns one ``h200x2`` 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. runs Module 1/2/3 across the dataset (per-camera VQA via PR 3471),
-  4. uploads the annotated dataset to ``--push_to_hub``.
-
-Usage:
-
-    HF_TOKEN=hf_... uv run python examples/annotation/run_hf_job.py
-
-Adjust ``CMD`` below to point at your own dataset / target hub repo.
-"""
-
-import os
-
-from huggingface_hub import get_token, run_job
-
-token = os.environ.get("HF_TOKEN") or get_token()
-if not token:
-    raise RuntimeError(
-        "No HF token. Run `huggingface-cli login` or `export HF_TOKEN=hf_...`"
-    )
-
-CMD = (
-    "apt-get update -qq && apt-get install -y -qq git ffmpeg && "
-    "pip install --no-deps "
-    "'lerobot @ git+https://github.com/huggingface/lerobot.git@feat/language-annotation-pipeline' && "
-    "pip install --upgrade-strategy only-if-needed "
-    "datasets pyarrow av jsonlines draccus gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect && "
-    "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
-    "export VLLM_VIDEO_BACKEND=pyav && "
-    "lerobot-annotate "
-    "--repo_id=imstevenpmwork/super_poulain_draft "
-    "--vlm.backend=openai "
-    "--vlm.model_id=Qwen/Qwen3.6-35B-A3B-FP8 "
-    "--vlm.parallel_servers=2 "
-    "--vlm.num_gpus=2 "
-    '--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}" '
-    "--vlm.serve_ready_timeout_s=1800 "
-    "--vlm.client_concurrency=256 "
-    "--vlm.max_new_tokens=512 "
-    "--executor.episode_parallelism=32 "
-    "--vlm.chat_template_kwargs='{enable_thinking: false}' "
-    "--vlm.camera_key=observation.images.wrist "
-    "--module_1.frames_per_second=1.0 "
-    "--module_1.use_video_url=true "
-    "--module_1.use_video_url_fps=1.0 "
-    "--module_3.K=1 --module_3.vqa_emission_hz=0.2 "
-    "--push_to_hub=pepijn223/super_poulain_qwen36moe-3"
-)
-
-job = run_job(
-    image="vllm/vllm-openai:latest",
-    command=["bash", "-c", CMD],
-    flavor="h200x2",
-    secrets={"HF_TOKEN": token},
-    timeout="2h",
-)
-print(f"Job URL: {job.url}")
-print(f"Job ID:  {job.id}")

examples/hil/hil_utils.py

@@ -1,226 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Shared utilities for Human-in-the-Loop data collection scripts."""
-
-import logging
-import time
-from dataclasses import dataclass, field
-from pathlib import Path
-
-from lerobot.common.control_utils import is_headless
-from lerobot.processor import (
-    IdentityProcessorStep,
-    RobotAction,
-    RobotObservation,
-    RobotProcessorPipeline,
-    observation_to_transition,
-    robot_action_observation_to_transition,
-    transition_to_observation,
-    transition_to_robot_action,
-)
-from lerobot.robots import Robot
-from lerobot.teleoperators import Teleoperator
-from lerobot.utils.robot_utils import precise_sleep
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class HILDatasetConfig:
-    repo_id: str
-    single_task: str
-    root: str | Path | None = None
-    fps: int = 30
-    episode_time_s: float = 120
-    num_episodes: int = 50
-    video: bool = True
-    push_to_hub: bool = True
-    private: bool = False
-    tags: list[str] | None = None
-    num_image_writer_processes: int = 0
-    num_image_writer_threads_per_camera: int = 4
-    video_encoding_batch_size: int = 1
-    vcodec: str = "auto"
-    streaming_encoding: bool = True
-    encoder_queue_maxsize: int = 30
-    encoder_threads: int | None = None
-    rename_map: dict[str, str] = field(default_factory=dict)
-
-
-def teleop_has_motor_control(teleop: Teleoperator) -> bool:
-    """Check if teleoperator has motor control capabilities."""
-    return all(hasattr(teleop, attr) for attr in ("enable_torque", "disable_torque", "write_goal_positions"))
-
-
-def teleop_disable_torque(teleop: Teleoperator) -> None:
-    """Disable teleop torque if supported."""
-    if hasattr(teleop, "disable_torque"):
-        teleop.disable_torque()
-
-
-def teleop_enable_torque(teleop: Teleoperator) -> None:
-    """Enable teleop torque if supported."""
-    if hasattr(teleop, "enable_torque"):
-        teleop.enable_torque()
-
-
-def teleop_smooth_move_to(teleop: Teleoperator, target_pos: dict, duration_s: float = 2.0, fps: int = 50):
-    """Smoothly move teleop to target position if motor control is available."""
-    if not teleop_has_motor_control(teleop):
-        logger.warning("Teleop does not support motor control - cannot mirror robot position")
-        return
-
-    teleop_enable_torque(teleop)
-    current = teleop.get_action()
-    steps = max(int(duration_s * fps), 1)
-
-    for step in range(steps + 1):
-        t = step / steps
-        interp = {}
-        for k in current:
-            if k in target_pos:
-                interp[k] = current[k] * (1 - t) + target_pos[k] * t
-            else:
-                interp[k] = current[k]
-        teleop.write_goal_positions(interp)
-        time.sleep(1 / fps)
-
-
-def init_keyboard_listener():
-    """Initialize keyboard listener with HIL controls."""
-    events = {
-        "exit_early": False,
-        "rerecord_episode": False,
-        "stop_recording": False,
-        "policy_paused": False,
-        "correction_active": False,
-        "resume_policy": False,
-        "in_reset": False,
-        "start_next_episode": False,
-    }
-
-    if is_headless():
-        logger.warning("Headless environment - keyboard controls unavailable")
-        return None, events
-
-    from pynput import keyboard
-
-    def on_press(key):
-        try:
-            if events["in_reset"]:
-                if key in [keyboard.Key.space, keyboard.Key.right]:
-                    logger.info("[HIL] Starting next episode...")
-                    events["start_next_episode"] = True
-                elif hasattr(key, "char") and key.char == "c":
-                    events["start_next_episode"] = True
-                elif key == keyboard.Key.esc:
-                    logger.info("[HIL] ESC - Stop recording, pushing to hub...")
-                    events["stop_recording"] = True
-                    events["start_next_episode"] = True
-            else:
-                if key == keyboard.Key.space:
-                    if not events["policy_paused"] and not events["correction_active"]:
-                        logger.info("[HIL] PAUSED - Press 'c' to take control or 'p' to resume policy")
-                        events["policy_paused"] = True
-                elif hasattr(key, "char") and key.char == "c":
-                    if events["policy_paused"] and not events["correction_active"]:
-                        logger.info("[HIL] Taking control...")
-                        events["start_next_episode"] = True
-                elif hasattr(key, "char") and key.char == "p":
-                    if events["policy_paused"] or events["correction_active"]:
-                        logger.info("[HIL] Resuming policy...")
-                        events["resume_policy"] = True
-                elif key == keyboard.Key.right:
-                    logger.info("[HIL] End episode")
-                    events["exit_early"] = True
-                elif key == keyboard.Key.left:
-                    logger.info("[HIL] Re-record episode")
-                    events["rerecord_episode"] = True
-                    events["exit_early"] = True
-                elif key == keyboard.Key.esc:
-                    logger.info("[HIL] ESC - Stop recording...")
-                    events["stop_recording"] = True
-                    events["exit_early"] = True
-        except Exception as e:
-            logger.info(f"Key error: {e}")
-
-    listener = keyboard.Listener(on_press=on_press)
-    listener.start()
-    return listener, events
-
-
-def make_identity_processors():
-    """Create identity processors for recording."""
-    teleop_proc = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-        steps=[IdentityProcessorStep()],
-        to_transition=robot_action_observation_to_transition,
-        to_output=transition_to_robot_action,
-    )
-    obs_proc = RobotProcessorPipeline[RobotObservation, RobotObservation](
-        steps=[IdentityProcessorStep()],
-        to_transition=observation_to_transition,
-        to_output=transition_to_observation,
-    )
-    return teleop_proc, obs_proc
-
-
-def reset_loop(robot: Robot, teleop: Teleoperator, events: dict, fps: int):
-    """Reset period where human repositions environment."""
-    logger.info("[HIL] RESET")
-
-    events["in_reset"] = True
-    events["start_next_episode"] = False
-
-    obs = robot.get_observation()
-    robot_pos = {k: v for k, v in obs.items() if k.endswith(".pos") and k in robot.observation_features}
-    teleop_smooth_move_to(teleop, robot_pos, duration_s=2.0, fps=50)
-
-    logger.info("Press any key to enable teleoperation")
-    while not events["start_next_episode"] and not events["stop_recording"]:
-        precise_sleep(0.05)
-
-    if events["stop_recording"]:
-        return
-
-    events["start_next_episode"] = False
-    teleop_disable_torque(teleop)
-    logger.info("Teleop enabled - press any key to start episode")
-
-    while not events["start_next_episode"] and not events["stop_recording"]:
-        loop_start = time.perf_counter()
-        action = teleop.get_action()
-        robot.send_action(action)
-        precise_sleep(1 / fps - (time.perf_counter() - loop_start))
-
-    events["in_reset"] = False
-    events["start_next_episode"] = False
-    events["exit_early"] = False
-    events["policy_paused"] = False
-    events["correction_active"] = False
-    events["resume_policy"] = False
-
-
-def print_controls(rtc: bool = False):
-    """Print control instructions."""
-    mode = "Human-in-the-Loop Data Collection" + (" (RTC)" if rtc else "")
-    logger.info(
-        "%s\n  Controls:\n"
-        "    SPACE  - Pause policy\n"
-        "    c      - Take control\n"
-        "    p      - Resume policy after pause/correction\n"
-        "    →      - End episode\n"
-        "    ESC    - Stop and push to hub",
-        mode,
-    )

examples/lekiwi/evaluate.py

@@ -14,17 +14,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from lerobot.common.control_utils import init_keyboard_listener
+import logging
+import time
+
+from lerobot.common.control_utils import init_keyboard_listener, predict_action
 from lerobot.datasets import LeRobotDataset
 from lerobot.policies import make_pre_post_processors
 from lerobot.policies.act import ACTPolicy
+from lerobot.policies.utils import make_robot_action
 from lerobot.processor import make_default_processors
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.utils.constants import ACTION, OBS_STR
-from lerobot.utils.feature_utils import hw_to_dataset_features
+from lerobot.utils.feature_utils import build_dataset_frame, hw_to_dataset_features
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
 NUM_EPISODES = 2
 FPS = 30
@@ -35,6 +39,9 @@ HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>"
 
 
 def main():
+    # NOTE: For production policy deployment, use `lerobot-rollout` CLI instead.
+    # This script provides a self-contained example for educational purposes.
+
     # Create the robot configuration & robot
     robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 
@@ -83,43 +90,67 @@ def main():
             raise ValueError("Robot is not connected!")
 
         print("Starting evaluate loop...")
+        control_interval = 1 / FPS
         recorded_episodes = 0
         while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
             log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}")
 
-            # Main record loop
-            record_loop(
-                robot=robot,
-                events=events,
-                fps=FPS,
-                policy=policy,
-                preprocessor=preprocessor,  # Pass the pre and post policy processors
-                postprocessor=postprocessor,
-                dataset=dataset,
-                control_time_s=EPISODE_TIME_SEC,
-                single_task=TASK_DESCRIPTION,
-                display_data=True,
-                teleop_action_processor=teleop_action_processor,
-                robot_action_processor=robot_action_processor,
-                robot_observation_processor=robot_observation_processor,
-            )
+            # Inline evaluation loop: predict actions and send to robot
+            timestamp = 0
+            start_episode_t = time.perf_counter()
+            while timestamp < EPISODE_TIME_SEC:
+                start_loop_t = time.perf_counter()
+
+                if events["exit_early"]:
+                    events["exit_early"] = False
+                    break
+
+                # Get robot observation
+                obs = robot.get_observation()
+                obs_processed = robot_observation_processor(obs)
+                observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix=OBS_STR)
+
+                # Predict action using the policy
+                action_tensor = predict_action(
+                    observation=observation_frame,
+                    policy=policy,
+                    device=policy.config.device,
+                    preprocessor=preprocessor,
+                    postprocessor=postprocessor,
+                    use_amp=policy.config.device.type == "cuda",
+                    task=TASK_DESCRIPTION,
+                    robot_type=robot.name,
+                )
+
+                # Convert policy output to robot action dict
+                action_values = make_robot_action(action_tensor, dataset.features)
+
+                # Process and send action to robot
+                robot_action_to_send = robot_action_processor((action_values, obs))
+                robot.send_action(robot_action_to_send)
+
+                # Write to dataset
+                action_frame = build_dataset_frame(dataset.features, action_values, prefix=ACTION)
+                frame = {**observation_frame, **action_frame, "task": TASK_DESCRIPTION}
+                dataset.add_frame(frame)
+
+                log_rerun_data(observation=obs_processed, action=action_values)
+
+                dt_s = time.perf_counter() - start_loop_t
+                sleep_time_s = control_interval - dt_s
+                if sleep_time_s < 0:
+                    logging.warning(
+                        f"Evaluate loop is running slower ({1 / dt_s:.1f} Hz) than the target FPS ({FPS} Hz)."
+                    )
+                precise_sleep(max(sleep_time_s, 0.0))
+                timestamp = time.perf_counter() - start_episode_t
 
             # Reset the environment if not stopping or re-recording
             if not events["stop_recording"] and (
                 (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
             ):
                 log_say("Reset the environment")
-                record_loop(
-                    robot=robot,
-                    events=events,
-                    fps=FPS,
-                    control_time_s=EPISODE_TIME_SEC,
-                    single_task=TASK_DESCRIPTION,
-                    display_data=True,
-                    teleop_action_processor=teleop_action_processor,
-                    robot_action_processor=robot_action_processor,
-                    robot_observation_processor=robot_observation_processor,
-                )
+                log_say("Waiting for environment reset, press right arrow key when ready...")
 
             if events["rerecord_episode"]:
                 log_say("Re-record episode")

examples/lekiwi/record.py

@@ -45,9 +45,6 @@ def main():
     leader_arm = SO100Leader(leader_arm_config)
     keyboard = KeyboardTeleop(keyboard_config)
 
-    # TODO(Steven): Update this example to use pipelines
-    teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
-
     # Configure the dataset features
     action_features = hw_to_dataset_features(robot.action_features, ACTION)
     obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
@@ -77,6 +74,10 @@ def main():
         if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
             raise ValueError("Robot or teleop is not connected!")
 
+        teleop_action_processor, robot_action_processor, robot_observation_processor = (
+            make_default_processors()
+        )
+
         print("Starting record loop...")
         recorded_episodes = 0
         while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
@@ -87,14 +88,14 @@ def main():
                 robot=robot,
                 events=events,
                 fps=FPS,
+                teleop_action_processor=teleop_action_processor,
+                robot_action_processor=robot_action_processor,
+                robot_observation_processor=robot_observation_processor,
                 dataset=dataset,
                 teleop=[leader_arm, keyboard],
                 control_time_s=EPISODE_TIME_SEC,
                 single_task=TASK_DESCRIPTION,
                 display_data=True,
-                teleop_action_processor=teleop_action_processor,
-                robot_action_processor=robot_action_processor,
-                robot_observation_processor=robot_observation_processor,
             )
 
             # Reset the environment if not stopping or re-recording
@@ -106,13 +107,13 @@ def main():
                     robot=robot,
                     events=events,
                     fps=FPS,
+                    teleop_action_processor=teleop_action_processor,
+                    robot_action_processor=robot_action_processor,
+                    robot_observation_processor=robot_observation_processor,
                     teleop=[leader_arm, keyboard],
                     control_time_s=RESET_TIME_SEC,
                     single_task=TASK_DESCRIPTION,
                     display_data=True,
-                    teleop_action_processor=teleop_action_processor,
-                    robot_action_processor=robot_action_processor,
-                    robot_observation_processor=robot_observation_processor,
                 )
 
             if events["rerecord_episode"]:

examples/lekiwi/rollout.py

@@ -0,0 +1,77 @@
+# !/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Run a trained policy on LeKiwi without recording (base rollout).
+
+Uses the rollout engine's :class:`BaseStrategy` (autonomous execution,
+no dataset) with :class:`SyncInferenceConfig` (inline policy call per
+control tick).  For a CLI entry point with the same capabilities plus
+recording, upload, and human-in-the-loop variants, see ``lerobot-rollout``.
+"""
+
+from lerobot.configs import PreTrainedConfig
+from lerobot.robots.lekiwi import LeKiwiClientConfig
+from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
+from lerobot.rollout.inference import SyncInferenceConfig
+from lerobot.rollout.strategies import BaseStrategy
+from lerobot.utils.process import ProcessSignalHandler
+from lerobot.utils.utils import init_logging
+
+FPS = 30
+DURATION_SEC = 60
+TASK_DESCRIPTION = "My task description"
+HF_MODEL_ID = "<hf_username>/<model_repo_id>"
+
+
+def main():
+    init_logging()
+
+    # Robot: LeKiwi client — make sure lekiwi_host is already running on the robot.
+    robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
+
+    # Policy: load the pretrained config.  ``pretrained_path`` is read downstream
+    # by ``build_rollout_context`` to reload the full model.
+    policy_config = PreTrainedConfig.from_pretrained(HF_MODEL_ID)
+    policy_config.pretrained_path = HF_MODEL_ID
+
+    # Assemble the rollout config: base strategy (no recording) + sync inference.
+    cfg = RolloutConfig(
+        robot=robot_config,
+        policy=policy_config,
+        strategy=BaseStrategyConfig(),
+        inference=SyncInferenceConfig(),
+        fps=FPS,
+        duration=DURATION_SEC,
+        task=TASK_DESCRIPTION,
+    )
+
+    # Graceful Ctrl-C: the strategy loop exits when shutdown_event is set.
+    signal_handler = ProcessSignalHandler(use_threads=True)
+
+    # Build the context (connects robot, loads policy, wires the inference strategy).
+    # No custom processors here — LeKiwi runs on raw joint features.
+    ctx = build_rollout_context(cfg, signal_handler.shutdown_event)
+
+    strategy = BaseStrategy(cfg.strategy)
+    try:
+        strategy.setup(ctx)
+        strategy.run(ctx)
+    finally:
+        strategy.teardown(ctx)
+
+
+if __name__ == "__main__":
+    main()

examples/phone_to_so100/evaluate.py

@@ -14,13 +14,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+import time
+
 from lerobot.cameras.opencv import OpenCVCameraConfig
-from lerobot.common.control_utils import init_keyboard_listener
+from lerobot.common.control_utils import init_keyboard_listener, predict_action
 from lerobot.configs import FeatureType, PolicyFeature
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.policies import make_pre_post_processors
 from lerobot.policies.act import ACTPolicy
+from lerobot.policies.utils import make_robot_action
 from lerobot.processor import (
     RobotProcessorPipeline,
     make_default_teleop_action_processor,
@@ -34,11 +38,12 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.types import RobotAction, RobotObservation
-from lerobot.utils.feature_utils import combine_feature_dicts
+from lerobot.utils.constants import ACTION, OBS_STR
+from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
 NUM_EPISODES = 5
 FPS = 30
@@ -49,6 +54,9 @@ HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
 
 
 def main():
+    # NOTE: For production policy deployment, use `lerobot-rollout` CLI instead.
+    # This script provides a self-contained example for educational purposes.
+
     # Create the robot configuration & robot
     camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
     robot_config = SO100FollowerConfig(
@@ -143,43 +151,67 @@ def main():
             raise ValueError("Robot is not connected!")
 
         print("Starting evaluate loop...")
+        control_interval = 1 / FPS
         episode_idx = 0
         for episode_idx in range(NUM_EPISODES):
             log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
 
-            # Main record loop
-            record_loop(
-                robot=robot,
-                events=events,
-                fps=FPS,
-                policy=policy,
-                preprocessor=preprocessor,  # Pass the pre and post policy processors
-                postprocessor=postprocessor,
-                dataset=dataset,
-                control_time_s=EPISODE_TIME_SEC,
-                single_task=TASK_DESCRIPTION,
-                display_data=True,
-                teleop_action_processor=make_default_teleop_action_processor(),
-                robot_action_processor=robot_ee_to_joints_processor,
-                robot_observation_processor=robot_joints_to_ee_pose_processor,
-            )
+            # Inline evaluation loop: predict actions and send to robot
+            timestamp = 0
+            start_episode_t = time.perf_counter()
+            while timestamp < EPISODE_TIME_SEC:
+                start_loop_t = time.perf_counter()
+
+                if events["exit_early"]:
+                    events["exit_early"] = False
+                    break
+
+                # Get robot observation
+                obs = robot.get_observation()
+                obs_processed = robot_joints_to_ee_pose_processor(obs)
+                observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix=OBS_STR)
+
+                # Predict action using the policy
+                action_tensor = predict_action(
+                    observation=observation_frame,
+                    policy=policy,
+                    device=policy.config.device,
+                    preprocessor=preprocessor,
+                    postprocessor=postprocessor,
+                    use_amp=policy.config.device.type == "cuda",
+                    task=TASK_DESCRIPTION,
+                    robot_type=robot.name,
+                )
+
+                # Convert policy output to robot action dict
+                action_values = make_robot_action(action_tensor, dataset.features)
+
+                # Process and send action to robot (EE -> joints via IK)
+                robot_action_to_send = robot_ee_to_joints_processor((action_values, obs))
+                robot.send_action(robot_action_to_send)
+
+                # Write to dataset
+                action_frame = build_dataset_frame(dataset.features, action_values, prefix=ACTION)
+                frame = {**observation_frame, **action_frame, "task": TASK_DESCRIPTION}
+                dataset.add_frame(frame)
+
+                log_rerun_data(observation=obs_processed, action=action_values)
+
+                dt_s = time.perf_counter() - start_loop_t
+                sleep_time_s = control_interval - dt_s
+                if sleep_time_s < 0:
+                    logging.warning(
+                        f"Evaluate loop is running slower ({1 / dt_s:.1f} Hz) than the target FPS ({FPS} Hz)."
+                    )
+                precise_sleep(max(sleep_time_s, 0.0))
+                timestamp = time.perf_counter() - start_episode_t
 
             # Reset the environment if not stopping or re-recording
             if not events["stop_recording"] and (
                 (episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]
             ):
                 log_say("Reset the environment")
-                record_loop(
-                    robot=robot,
-                    events=events,
-                    fps=FPS,
-                    control_time_s=EPISODE_TIME_SEC,
-                    single_task=TASK_DESCRIPTION,
-                    display_data=True,
-                    teleop_action_processor=make_default_teleop_action_processor(),
-                    robot_action_processor=robot_ee_to_joints_processor,
-                    robot_observation_processor=robot_joints_to_ee_pose_processor,
-                )
+                log_say("Waiting for environment reset, press right arrow key when ready...")
 
             if events["rerecord_episode"]:
                 log_say("Re-record episode")
@@ -190,7 +222,6 @@ def main():
 
             # Save episode
             dataset.save_episode()
-            episode_idx += 1
     finally:
         # Clean up
         log_say("Stop recording")

examples/phone_to_so100/record.py

@@ -65,14 +65,15 @@ def main():
     robot = SO100Follower(robot_config)
     phone = Phone(teleop_config)
 
-    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo:
+    #   https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
     kinematics_solver = RobotKinematics(
         urdf_path="./SO101/so101_new_calib.urdf",
         target_frame_name="gripper_frame_link",
         joint_names=list(robot.bus.motors.keys()),
     )
 
-    # Build pipeline to convert phone action to EE action
+    # Build pipeline to convert phone action to EE action (with gripper velocity mapped to joint).
     phone_to_robot_ee_pose_processor = RobotProcessorPipeline[
         tuple[RobotAction, RobotObservation], RobotAction
     ](
@@ -94,7 +95,7 @@ def main():
         to_output=transition_to_robot_action,
     )
 
-    # Build pipeline to convert EE action to joints action
+    # Build pipeline to convert EE action to joints action (IK).
     robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
         steps=[
             InverseKinematicsEEToJoints(
@@ -107,7 +108,7 @@ def main():
         to_output=transition_to_robot_action,
     )
 
-    # Build pipeline to convert joint observation to EE observation
+    # Build pipeline to convert joint observation to EE observation (FK).
     robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation](
         steps=[
             ForwardKinematicsJointsToEE(
@@ -118,13 +119,12 @@ def main():
         to_output=transition_to_observation,
     )
 
-    # Create the dataset
+    # Create the dataset, deriving features from the pipelines so the on-disk schema
+    # matches exactly what the pipelines produce at runtime.
     dataset = LeRobotDataset.create(
         repo_id=HF_REPO_ID,
         fps=FPS,
         features=combine_feature_dicts(
-            # Run the feature contract of the pipelines
-            # This tells you how the features would look like after the pipeline steps
             aggregate_pipeline_dataset_features(
                 pipeline=phone_to_robot_ee_pose_processor,
                 initial_features=create_initial_features(action=phone.action_features),
@@ -163,14 +163,14 @@ def main():
                 robot=robot,
                 events=events,
                 fps=FPS,
+                teleop_action_processor=phone_to_robot_ee_pose_processor,
+                robot_action_processor=robot_ee_to_joints_processor,
+                robot_observation_processor=robot_joints_to_ee_pose,
                 teleop=phone,
                 dataset=dataset,
                 control_time_s=EPISODE_TIME_SEC,
                 single_task=TASK_DESCRIPTION,
                 display_data=True,
-                teleop_action_processor=phone_to_robot_ee_pose_processor,
-                robot_action_processor=robot_ee_to_joints_processor,
-                robot_observation_processor=robot_joints_to_ee_pose,
             )
 
             # Reset the environment if not stopping or re-recording
@@ -182,13 +182,13 @@ def main():
                     robot=robot,
                     events=events,
                     fps=FPS,
+                    teleop_action_processor=phone_to_robot_ee_pose_processor,
+                    robot_action_processor=robot_ee_to_joints_processor,
+                    robot_observation_processor=robot_joints_to_ee_pose,
                     teleop=phone,
                     control_time_s=RESET_TIME_SEC,
                     single_task=TASK_DESCRIPTION,
                     display_data=True,
-                    teleop_action_processor=phone_to_robot_ee_pose_processor,
-                    robot_action_processor=robot_ee_to_joints_processor,
-                    robot_observation_processor=robot_joints_to_ee_pose,
                 )
 
             if events["rerecord_episode"]:

examples/phone_to_so100/rollout.py

@@ -0,0 +1,126 @@
+# !/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Run a trained EE-space policy on SO100 (phone-trained) without recording.
+
+Mirrors ``examples/so100_to_so100_EE/rollout.py`` — the model was trained
+with phone teleoperation in EE space, so at deployment we only need the
+joint↔EE conversion on the robot side; the phone is not used.
+
+Uses :class:`BaseStrategy` (no recording) + :class:`SyncInferenceConfig`
+(inline policy call).  For recording during rollout, switch to Sentry,
+Highlight, or DAgger via ``lerobot-rollout --strategy.type=...``.
+"""
+
+from lerobot.cameras.opencv import OpenCVCameraConfig
+from lerobot.configs import PreTrainedConfig
+from lerobot.model.kinematics import RobotKinematics
+from lerobot.processor import (
+    RobotProcessorPipeline,
+    observation_to_transition,
+    robot_action_observation_to_transition,
+    transition_to_observation,
+    transition_to_robot_action,
+)
+from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
+from lerobot.robots.so_follower.robot_kinematic_processor import (
+    ForwardKinematicsJointsToEE,
+    InverseKinematicsEEToJoints,
+)
+from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
+from lerobot.rollout.inference import SyncInferenceConfig
+from lerobot.rollout.strategies import BaseStrategy
+from lerobot.types import RobotAction, RobotObservation
+from lerobot.utils.process import ProcessSignalHandler
+from lerobot.utils.utils import init_logging
+
+FPS = 30
+DURATION_SEC = 60
+TASK_DESCRIPTION = "My task description"
+HF_MODEL_ID = "<hf_username>/<model_repo_id>"
+
+
+def main():
+    init_logging()
+
+    camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
+    robot_config = SO100FollowerConfig(
+        port="/dev/tty.usbmodem58760434471",
+        id="my_awesome_follower_arm",
+        cameras=camera_config,
+        use_degrees=True,
+    )
+
+    # Peek at motor names once to build the kinematic solver.
+    temp_robot = SO100Follower(robot_config)
+    motor_names = list(temp_robot.bus.motors.keys())
+
+    kinematics_solver = RobotKinematics(
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=motor_names,
+    )
+
+    robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
+        steps=[ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=motor_names)],
+        to_transition=observation_to_transition,
+        to_output=transition_to_observation,
+    )
+
+    robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+        steps=[
+            InverseKinematicsEEToJoints(
+                kinematics=kinematics_solver,
+                motor_names=motor_names,
+                initial_guess_current_joints=True,
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
+
+    policy_config = PreTrainedConfig.from_pretrained(HF_MODEL_ID)
+    policy_config.pretrained_path = HF_MODEL_ID
+
+    cfg = RolloutConfig(
+        robot=robot_config,
+        policy=policy_config,
+        strategy=BaseStrategyConfig(),
+        inference=SyncInferenceConfig(),
+        fps=FPS,
+        duration=DURATION_SEC,
+        task=TASK_DESCRIPTION,
+    )
+
+    signal_handler = ProcessSignalHandler(use_threads=True)
+
+    ctx = build_rollout_context(
+        cfg,
+        signal_handler.shutdown_event,
+        robot_action_processor=robot_ee_to_joints_processor,
+        robot_observation_processor=robot_joints_to_ee_pose_processor,
+    )
+
+    strategy = BaseStrategy(cfg.strategy)
+    try:
+        strategy.setup(ctx)
+        strategy.run(ctx)
+    finally:
+        strategy.teardown(ctx)
+
+
+if __name__ == "__main__":
+    main()

examples/rtc/eval_with_real_robot.py

@@ -1,673 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Demo script showing how to use Real-Time Chunking (RTC) with action chunking policies on real robots.
-
-This script demonstrates:
-1. Creating a robot and policy (SmolVLA, Pi0, etc.) with RTC
-2. Consuming actions from the policy while the robot executes
-3. Periodically requesting new action chunks in the background using threads
-4. Managing action buffers and timing for real-time operation
-
-For simulation environments, see eval_with_simulation.py
-
-Usage:
-    # Run RTC with Real robot with RTC
-    uv run examples/rtc/eval_with_real_robot.py \
-        --policy.path=<USER>/smolvla_check_rtc_last3 \
-        --policy.device=mps \
-        --rtc.enabled=true \
-        --rtc.execution_horizon=20 \
-        --robot.type=so100_follower \
-        --robot.port=/dev/tty.usbmodem58FA0834591 \
-        --robot.id=so100_follower \
-        --robot.cameras="{ gripper: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
-        --task="Move green small object into the purple platform" \
-        --duration=120
-
-    # Run RTC with Real robot without RTC
-    uv run examples/rtc/eval_with_real_robot.py \
-        --policy.path=<USER>/smolvla_check_rtc_last3 \
-        --policy.device=mps \
-        --rtc.enabled=false \
-        --robot.type=so100_follower \
-        --robot.port=/dev/tty.usbmodem58FA0834591 \
-        --robot.id=so100_follower \
-        --robot.cameras="{ gripper: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
-        --task="Move green small object into the purple platform" \
-        --duration=120
-
-    # Run RTC with Real robot with pi0.5 policy
-    uv run examples/rtc/eval_with_real_robot.py \
-        --policy.path=<USER>/pi05_check_rtc \
-        --policy.device=mps \
-        --rtc.enabled=true \
-        --rtc.execution_horizon=20 \
-        --robot.type=so100_follower \
-        --robot.port=/dev/tty.usbmodem58FA0834591 \
-        --robot.id=so100_follower \
-        --robot.cameras="{ gripper: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}, front: {type: opencv, index_or_path: 1, width: 640, height: 480, fps: 30}}" \
-        --task="Move green small object into the purple platform" \
-        --duration=120
-
-    # Run RTC with bi_openarm_follower (dual-arm OpenArms) and pi0.5 policy
-    python examples/rtc/eval_with_real_robot.py \
-        --policy.path=lerobot-data-collection/folding_final \
-        --robot.type=bi_openarm_follower \
-        --robot.cameras='{left_wrist: {type: opencv, index_or_path: "/dev/video4", width: 1280, height: 720, fps: 30}, base: {type: opencv, index_or_path: "/dev/video2", width: 640, height: 480, fps: 30}, right_wrist: {type: opencv, index_or_path: "/dev/video0", width: 1280, height: 720, fps: 30}}' \
-        --robot.left_arm_config.port=can0 \
-        --robot.left_arm_config.side=left \
-        --robot.left_arm_config.can_interface=socketcan \
-        --robot.left_arm_config.disable_torque_on_disconnect=true \
-        --robot.left_arm_config.max_relative_target=8.0 \
-        --robot.right_arm_config.port=can1 \
-        --robot.right_arm_config.side=right \
-        --robot.right_arm_config.can_interface=socketcan \
-        --robot.right_arm_config.disable_torque_on_disconnect=true \
-        --robot.right_arm_config.max_relative_target=8.0 \
-        --task="Fold the T-shirt properly" \
-        --fps=30 \
-        --duration=2000 \
-        --interpolation_multiplier=3 \
-        --rtc.enabled=true \
-        --rtc.execution_horizon=20 \
-        --rtc.max_guidance_weight=5.0 \
-        --rtc.prefix_attention_schedule=LINEAR \
-        --device=cuda
-"""
-
-import logging
-import math
-import sys
-import time
-import traceback
-from dataclasses import dataclass, field
-from threading import Event, Lock, Thread
-
-import torch
-from torch import Tensor
-
-from lerobot.cameras.opencv import OpenCVCameraConfig  # noqa: F401
-from lerobot.cameras.realsense import RealSenseCameraConfig  # noqa: F401
-from lerobot.cameras.zmq import ZMQCameraConfig  # noqa: F401
-from lerobot.configs import PreTrainedConfig, RTCAttentionSchedule, parser
-from lerobot.policies import get_policy_class, make_pre_post_processors
-from lerobot.policies.rtc import ActionInterpolator, ActionQueue, LatencyTracker, RTCConfig
-from lerobot.processor import (
-    NormalizerProcessorStep,
-    RelativeActionsProcessorStep,
-    TransitionKey,
-    create_transition,
-    make_default_robot_action_processor,
-    make_default_robot_observation_processor,
-    to_relative_actions,
-)
-from lerobot.rl.process import ProcessSignalHandler
-from lerobot.robots import (  # noqa: F401
-    Robot,
-    RobotConfig,
-    bi_openarm_follower,
-    bi_so_follower,
-    koch_follower,
-    so_follower,
-    unitree_g1,
-)
-from lerobot.robots.utils import make_robot_from_config
-from lerobot.utils.constants import OBS_IMAGES, OBS_STATE
-from lerobot.utils.feature_utils import build_dataset_frame, hw_to_dataset_features
-from lerobot.utils.hub import HubMixin
-from lerobot.utils.utils import init_logging
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class RobotWrapper:
-    def __init__(self, robot: Robot):
-        self.robot = robot
-        self.lock = Lock()
-
-    def get_observation(self) -> dict[str, Tensor]:
-        with self.lock:
-            return self.robot.get_observation()
-
-    def send_action(self, action: Tensor):
-        with self.lock:
-            self.robot.send_action(action)
-
-    def observation_features(self) -> list[str]:
-        with self.lock:
-            return self.robot.observation_features
-
-    def action_features(self) -> list[str]:
-        with self.lock:
-            return self.robot.action_features
-
-
-@dataclass
-class RTCDemoConfig(HubMixin):
-    """Configuration for RTC demo with action chunking policies and real robots."""
-
-    # Policy configuration
-    policy: PreTrainedConfig | None = None
-
-    # Robot configuration
-    robot: RobotConfig | None = None
-
-    # RTC configuration
-    rtc: RTCConfig = field(
-        default_factory=lambda: RTCConfig(
-            execution_horizon=10,
-            max_guidance_weight=1.0,
-            prefix_attention_schedule=RTCAttentionSchedule.EXP,
-        )
-    )
-
-    # Demo parameters
-    duration: float = 30.0  # Duration to run the demo (seconds)
-    fps: float = 10.0  # Action execution frequency (Hz)
-    interpolation_multiplier: int = 1  # Control rate multiplier (1=off, 2=2x, 3=3x)
-
-    # Compute device
-    device: str | None = None  # Device to run on (cuda, cpu, auto)
-
-    # Get new actions horizon. The amount of executed steps after which will be requested new actions.
-    # It should be higher than inference delay + execution horizon.
-    action_queue_size_to_get_new_actions: int = 30
-
-    # Task to execute
-    task: str = field(default="", metadata={"help": "Task to execute"})
-
-    # Torch compile configuration
-    use_torch_compile: bool = field(
-        default=False,
-        metadata={"help": "Use torch.compile for faster inference (PyTorch 2.0+)"},
-    )
-
-    torch_compile_backend: str = field(
-        default="inductor",
-        metadata={"help": "Backend for torch.compile (inductor, aot_eager, cudagraphs)"},
-    )
-
-    torch_compile_mode: str = field(
-        default="default",
-        metadata={"help": "Compilation mode (default, reduce-overhead, max-autotune)"},
-    )
-
-    torch_compile_disable_cudagraphs: bool = field(
-        default=True,
-        metadata={
-            "help": "Disable CUDA graphs in torch.compile. Required due to in-place tensor "
-            "operations in denoising loop (x_t += dt * v_t) which cause tensor aliasing issues."
-        },
-    )
-
-    def __post_init__(self):
-        # HACK: We parse again the cli args here to get the pretrained path if there was one.
-        policy_path = parser.get_path_arg("policy")
-        if policy_path:
-            cli_overrides = parser.get_cli_overrides("policy")
-            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
-            self.policy.pretrained_path = policy_path
-        else:
-            raise ValueError("Policy path is required")
-
-        # Validate that robot configuration is provided
-        if self.robot is None:
-            raise ValueError("Robot configuration must be provided")
-
-    @classmethod
-    def __get_path_fields__(cls) -> list[str]:
-        """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
-        return ["policy"]
-
-
-def is_image_key(k: str) -> bool:
-    return k.startswith(OBS_IMAGES)
-
-
-def _reanchor_relative_rtc_prefix(
-    prev_actions_absolute: Tensor,
-    current_state: Tensor,
-    relative_step: RelativeActionsProcessorStep,
-    normalizer_step: NormalizerProcessorStep | None,
-    policy_device: torch.device | str,
-) -> Tensor:
-    """Convert absolute leftovers into model-space for relative-action RTC policies.
-
-    When a policy uses relative actions, the RTC prefix (leftover actions from
-    the previous chunk) is stored in absolute space. Before feeding it back to
-    the policy we need to re-express it relative to the *current* robot state
-    and then re-normalize.
-    """
-    state = current_state.detach().cpu()
-    if state.dim() == 1:
-        state = state.unsqueeze(0)
-
-    action_cpu = prev_actions_absolute.detach().cpu()
-    mask = relative_step._build_mask(action_cpu.shape[-1])
-    relative_actions = to_relative_actions(action_cpu, state, mask)
-
-    transition = create_transition(action=relative_actions)
-    if normalizer_step is not None:
-        transition = normalizer_step(transition)
-
-    return transition[TransitionKey.ACTION].to(policy_device)
-
-
-def get_actions(
-    policy,
-    robot: RobotWrapper,
-    robot_observation_processor,
-    action_queue: ActionQueue,
-    shutdown_event: Event,
-    cfg: RTCDemoConfig,
-):
-    """Thread function to request action chunks from the policy.
-
-    Args:
-        policy: The policy instance (SmolVLA, Pi0, etc.)
-        robot: The robot instance for getting observations
-        robot_observation_processor: Processor for raw robot observations
-        action_queue: Queue to put new action chunks
-        shutdown_event: Event to signal shutdown
-        cfg: Demo configuration
-    """
-    try:
-        logger.info("[GET_ACTIONS] Starting get actions thread")
-
-        latency_tracker = LatencyTracker()  # Track latency of action chunks
-        fps = cfg.fps
-        time_per_chunk = 1.0 / fps
-
-        # Only keep .pos joints + camera streams if the policy was trained on positions,
-        # not the full pos/vel/torque state the robot exposes.
-        observation_features_hw = {
-            key: value
-            for key, value in robot.observation_features().items()
-            if key.endswith(".pos") or isinstance(value, tuple)
-        }
-
-        dataset_features = hw_to_dataset_features(observation_features_hw, "observation")
-        policy_device = policy.config.device
-
-        # Load preprocessor and postprocessor from pretrained files
-        # The stats are embedded in the processor .safetensors files
-        logger.info(f"[GET_ACTIONS] Loading preprocessor/postprocessor from {cfg.policy.pretrained_path}")
-
-        preprocessor, postprocessor = make_pre_post_processors(
-            policy_cfg=cfg.policy,
-            pretrained_path=cfg.policy.pretrained_path,
-            dataset_stats=None,  # Will load from pretrained processor files
-            preprocessor_overrides={
-                "device_processor": {"device": cfg.policy.device},
-            },
-        )
-
-        logger.info("[GET_ACTIONS] Preprocessor/postprocessor loaded successfully with embedded stats")
-
-        relative_step = next(
-            (s for s in preprocessor.steps if isinstance(s, RelativeActionsProcessorStep) and s.enabled),
-            None,
-        )
-        normalizer_step = next(
-            (s for s in preprocessor.steps if isinstance(s, NormalizerProcessorStep)),
-            None,
-        )
-        if relative_step is not None:
-            if relative_step.action_names is None:
-                cfg_names = getattr(cfg.policy, "action_feature_names", None)
-                if cfg_names:
-                    relative_step.action_names = list(cfg_names)
-                else:
-                    relative_step.action_names = [
-                        k for k in robot.robot.action_features if k.endswith(".pos")
-                    ]
-            logger.info("[GET_ACTIONS] Relative actions enabled: will re-anchor RTC prefix")
-
-        get_actions_threshold = cfg.action_queue_size_to_get_new_actions
-
-        if not cfg.rtc.enabled:
-            get_actions_threshold = 0
-
-        while not shutdown_event.is_set():
-            if action_queue.qsize() <= get_actions_threshold:
-                current_time = time.perf_counter()
-                action_index_before_inference = action_queue.get_action_index()
-                prev_actions = action_queue.get_left_over()
-
-                inference_latency = latency_tracker.max()
-                inference_delay = math.ceil(inference_latency / time_per_chunk)
-
-                obs = robot.get_observation()
-
-                # Apply robot observation processor
-                obs_processed = robot_observation_processor(obs)
-
-                obs_with_policy_features = build_dataset_frame(
-                    dataset_features, obs_processed, prefix="observation"
-                )
-
-                for name in obs_with_policy_features:
-                    obs_with_policy_features[name] = torch.from_numpy(obs_with_policy_features[name])
-                    if "image" in name:
-                        obs_with_policy_features[name] = (
-                            obs_with_policy_features[name].type(torch.float32) / 255
-                        )
-                        obs_with_policy_features[name] = (
-                            obs_with_policy_features[name].permute(2, 0, 1).contiguous()
-                        )
-                    obs_with_policy_features[name] = obs_with_policy_features[name].unsqueeze(0)
-                    obs_with_policy_features[name] = obs_with_policy_features[name].to(policy_device)
-
-                obs_with_policy_features["task"] = [cfg.task]  # Task should be a list, not a string!
-                obs_with_policy_features["robot_type"] = (
-                    robot.robot.name if hasattr(robot.robot, "name") else ""
-                )
-
-                preproceseded_obs = preprocessor(obs_with_policy_features)
-
-                # Re-anchor leftover actions for relative-action policies.
-                # We need the *postprocessed* (absolute) leftover, not the original
-                # (normalized/relative) one that get_left_over() returns.
-                if (
-                    prev_actions is not None
-                    and relative_step is not None
-                    and OBS_STATE in obs_with_policy_features
-                ):
-                    with action_queue.lock:
-                        if action_queue.queue is not None:
-                            prev_actions_abs = action_queue.queue[action_queue.last_index :].clone()
-                        else:
-                            prev_actions_abs = None
-                    if prev_actions_abs is not None and prev_actions_abs.numel() > 0:
-                        prev_actions = _reanchor_relative_rtc_prefix(
-                            prev_actions_absolute=prev_actions_abs,
-                            current_state=obs_with_policy_features[OBS_STATE],
-                            relative_step=relative_step,
-                            normalizer_step=normalizer_step,
-                            policy_device=policy_device,
-                        )
-
-                # Generate actions WITH RTC
-                actions = policy.predict_action_chunk(
-                    preproceseded_obs,
-                    inference_delay=inference_delay,
-                    prev_chunk_left_over=prev_actions,
-                )
-
-                # Store original actions (before postprocessing) for RTC
-                original_actions = actions.squeeze(0).clone()
-
-                postprocessed_actions = postprocessor(actions)
-
-                postprocessed_actions = postprocessed_actions.squeeze(0)
-
-                new_latency = time.perf_counter() - current_time
-                new_delay = math.ceil(new_latency / time_per_chunk)
-                latency_tracker.add(new_latency)
-
-                if cfg.action_queue_size_to_get_new_actions < cfg.rtc.execution_horizon + new_delay:
-                    logger.warning(
-                        "[GET_ACTIONS] cfg.action_queue_size_to_get_new_actions Too small, It should be higher than inference delay + execution horizon."
-                    )
-
-                action_queue.merge(
-                    original_actions, postprocessed_actions, new_delay, action_index_before_inference
-                )
-            else:
-                # Small sleep to prevent busy waiting
-                time.sleep(0.1)
-
-        logger.info("[GET_ACTIONS] get actions thread shutting down")
-    except Exception as e:
-        logger.error(f"[GET_ACTIONS] Fatal exception in get_actions thread: {e}")
-        logger.error(traceback.format_exc())
-        sys.exit(1)
-
-
-def actor_control(
-    robot: RobotWrapper,
-    robot_action_processor,
-    action_queue: ActionQueue,
-    shutdown_event: Event,
-    cfg: RTCDemoConfig,
-):
-    """Thread function to execute actions on the robot.
-
-    Args:
-        robot: The robot instance
-        action_queue: Queue to get actions from
-        shutdown_event: Event to signal shutdown
-        cfg: Demo configuration
-    """
-    try:
-        logger.info("[ACTOR] Starting actor thread")
-
-        action_keys = [k for k in robot.action_features() if k.endswith(".pos")]
-
-        action_count = 0
-        interpolator = ActionInterpolator(multiplier=cfg.interpolation_multiplier)
-        action_interval = interpolator.get_control_interval(cfg.fps)
-
-        while not shutdown_event.is_set():
-            start_time = time.perf_counter()
-
-            if interpolator.needs_new_action():
-                new_action = action_queue.get()
-                if new_action is not None:
-                    interpolator.add(new_action.cpu())
-
-            action = interpolator.get()
-            if action is not None:
-                action = action.cpu()
-                action_dict = {key: action[i].item() for i, key in enumerate(action_keys)}
-                action_processed = robot_action_processor((action_dict, None))
-                robot.send_action(action_processed)
-                action_count += 1
-
-            dt_s = time.perf_counter() - start_time
-            time.sleep(max(0, (action_interval - dt_s) - 0.001))
-
-        logger.info(f"[ACTOR] Actor thread shutting down. Total actions executed: {action_count}")
-    except Exception as e:
-        logger.error(f"[ACTOR] Fatal exception in actor_control thread: {e}")
-        logger.error(traceback.format_exc())
-        sys.exit(1)
-
-
-def _apply_torch_compile(policy, cfg: RTCDemoConfig):
-    """Apply torch.compile to the policy's predict_action_chunk method.
-
-    Args:
-        policy: Policy instance to compile
-        cfg: Configuration containing torch compile settings
-
-    Returns:
-        Policy with compiled predict_action_chunk method
-    """
-
-    # PI models handle their own compilation
-    if policy.type == "pi05" or policy.type == "pi0":
-        return policy
-
-    try:
-        # Check if torch.compile is available (PyTorch 2.0+)
-        if not hasattr(torch, "compile"):
-            logger.warning(
-                f"torch.compile is not available. Requires PyTorch 2.0+. "
-                f"Current version: {torch.__version__}. Skipping compilation."
-            )
-            return policy
-
-        logger.info("Applying torch.compile to predict_action_chunk...")
-        logger.info(f"  Backend: {cfg.torch_compile_backend}")
-        logger.info(f"  Mode: {cfg.torch_compile_mode}")
-        logger.info(f"  Disable CUDA graphs: {cfg.torch_compile_disable_cudagraphs}")
-
-        # Compile the predict_action_chunk method
-        # - CUDA graphs disabled to prevent tensor aliasing from in-place ops (x_t += dt * v_t)
-        compile_kwargs = {
-            "backend": cfg.torch_compile_backend,
-            "mode": cfg.torch_compile_mode,
-        }
-
-        # Disable CUDA graphs if requested (prevents tensor aliasing issues)
-        if cfg.torch_compile_disable_cudagraphs:
-            compile_kwargs["options"] = {"triton.cudagraphs": False}
-
-        original_method = policy.predict_action_chunk
-        compiled_method = torch.compile(original_method, **compile_kwargs)
-        policy.predict_action_chunk = compiled_method
-        logger.info("✓ Successfully compiled predict_action_chunk")
-
-    except Exception as e:
-        logger.error(f"Failed to apply torch.compile: {e}")
-        logger.warning("Continuing without torch.compile")
-
-    return policy
-
-
-@parser.wrap()
-def demo_cli(cfg: RTCDemoConfig):
-    """Main entry point for RTC demo with draccus configuration."""
-
-    # Initialize logging
-    init_logging()
-
-    logger.info(f"Using device: {cfg.device}")
-
-    # Setup signal handler for graceful shutdown
-    signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False)
-    shutdown_event = signal_handler.shutdown_event
-
-    policy = None
-    robot = None
-    get_actions_thread = None
-    actor_thread = None
-
-    policy_class = get_policy_class(cfg.policy.type)
-
-    # Load config and set compile_model for pi0/pi05 models
-    config = PreTrainedConfig.from_pretrained(cfg.policy.pretrained_path)
-
-    if cfg.policy.type == "pi05" or cfg.policy.type == "pi0":
-        config.compile_model = cfg.use_torch_compile
-
-    if config.use_peft:
-        from peft import PeftConfig, PeftModel
-
-        peft_pretrained_path = cfg.policy.pretrained_path
-        peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
-
-        policy = policy_class.from_pretrained(
-            pretrained_name_or_path=peft_config.base_model_name_or_path, config=config
-        )
-        policy = PeftModel.from_pretrained(policy, peft_pretrained_path, config=peft_config)
-    else:
-        policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=config)
-
-    # Turn on RTC
-    policy.config.rtc_config = cfg.rtc
-
-    # Init RTC processort, as by default if RTC disabled in the config
-    # The processor won't be created
-    policy.init_rtc_processor()
-
-    assert policy.name in ["smolvla", "pi05", "pi0"], "Only smolvla, pi05, and pi0 are supported for RTC"
-
-    policy = policy.to(cfg.device)
-    policy.eval()
-
-    # Apply torch.compile to predict_action_chunk method if enabled
-    if cfg.use_torch_compile:
-        policy = _apply_torch_compile(policy, cfg)
-
-    # Create robot
-    logger.info(f"Initializing robot: {cfg.robot.type}")
-    robot = make_robot_from_config(cfg.robot)
-    robot.connect()
-    robot_wrapper = RobotWrapper(robot)
-
-    # Create robot observation processor
-    robot_observation_processor = make_default_robot_observation_processor()
-    robot_action_processor = make_default_robot_action_processor()
-
-    # Create action queue for communication between threads
-    action_queue = ActionQueue(cfg.rtc)
-
-    # Start chunk requester thread
-    get_actions_thread = Thread(
-        target=get_actions,
-        args=(policy, robot_wrapper, robot_observation_processor, action_queue, shutdown_event, cfg),
-        daemon=True,
-        name="GetActions",
-    )
-    get_actions_thread.start()
-    logger.info("Started get actions thread")
-
-    # Start action executor thread
-    actor_thread = Thread(
-        target=actor_control,
-        args=(robot_wrapper, robot_action_processor, action_queue, shutdown_event, cfg),
-        daemon=True,
-        name="Actor",
-    )
-    actor_thread.start()
-    logger.info("Started actor thread")
-
-    logger.info("Started stop by duration thread")
-
-    # Main thread monitors for duration or shutdown
-    logger.info(f"Running demo for {cfg.duration} seconds...")
-    start_time = time.time()
-
-    while not shutdown_event.is_set() and (time.time() - start_time) < cfg.duration:
-        time.sleep(10)
-
-        # Log queue status periodically
-        if int(time.time() - start_time) % 5 == 0:
-            logger.info(f"[MAIN] Action queue size: {action_queue.qsize()}")
-
-        if time.time() - start_time > cfg.duration:
-            break
-
-    logger.info("Demo duration reached or shutdown requested")
-
-    # Signal shutdown
-    shutdown_event.set()
-
-    # Wait for threads to finish
-    if get_actions_thread and get_actions_thread.is_alive():
-        logger.info("Waiting for chunk requester thread to finish...")
-        get_actions_thread.join()
-
-    if actor_thread and actor_thread.is_alive():
-        logger.info("Waiting for action executor thread to finish...")
-        actor_thread.join()
-
-    # Cleanup robot
-    if robot:
-        robot.disconnect()
-        logger.info("Robot disconnected")
-
-    logger.info("Cleanup completed")
-
-
-if __name__ == "__main__":
-    demo_cli()
-    logging.info("RTC demo finished")

examples/so100_to_so100_EE/evaluate.py

@@ -14,13 +14,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
+import time
+
 from lerobot.cameras.opencv import OpenCVCameraConfig
-from lerobot.common.control_utils import init_keyboard_listener
+from lerobot.common.control_utils import init_keyboard_listener, predict_action
 from lerobot.configs import FeatureType, PolicyFeature
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.policies import make_pre_post_processors
 from lerobot.policies.act import ACTPolicy
+from lerobot.policies.utils import make_robot_action
 from lerobot.processor import (
     RobotProcessorPipeline,
     make_default_teleop_action_processor,
@@ -34,11 +38,12 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.types import RobotAction, RobotObservation
-from lerobot.utils.feature_utils import combine_feature_dicts
+from lerobot.utils.constants import ACTION, OBS_STR
+from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
+from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 
 NUM_EPISODES = 5
 FPS = 30
@@ -49,6 +54,9 @@ HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
 
 
 def main():
+    # NOTE: For production policy deployment, use `lerobot-rollout` CLI instead.
+    # This script provides a self-contained example for educational purposes.
+
     # Create the robot configuration & robot
     camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
     robot_config = SO100FollowerConfig(
@@ -143,43 +151,67 @@ def main():
             raise ValueError("Robot is not connected!")
 
         print("Starting evaluate loop...")
+        control_interval = 1 / FPS
         episode_idx = 0
         for episode_idx in range(NUM_EPISODES):
             log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
 
-            # Main record loop
-            record_loop(
-                robot=robot,
-                events=events,
-                fps=FPS,
-                policy=policy,
-                preprocessor=preprocessor,  # Pass the pre and post policy processors
-                postprocessor=postprocessor,
-                dataset=dataset,
-                control_time_s=EPISODE_TIME_SEC,
-                single_task=TASK_DESCRIPTION,
-                display_data=True,
-                teleop_action_processor=make_default_teleop_action_processor(),
-                robot_action_processor=robot_ee_to_joints_processor,
-                robot_observation_processor=robot_joints_to_ee_pose_processor,
-            )
+            # Inline evaluation loop: predict actions and send to robot
+            timestamp = 0
+            start_episode_t = time.perf_counter()
+            while timestamp < EPISODE_TIME_SEC:
+                start_loop_t = time.perf_counter()
+
+                if events["exit_early"]:
+                    events["exit_early"] = False
+                    break
+
+                # Get robot observation
+                obs = robot.get_observation()
+                obs_processed = robot_joints_to_ee_pose_processor(obs)
+                observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix=OBS_STR)
+
+                # Predict action using the policy
+                action_tensor = predict_action(
+                    observation=observation_frame,
+                    policy=policy,
+                    device=policy.config.device,
+                    preprocessor=preprocessor,
+                    postprocessor=postprocessor,
+                    use_amp=policy.config.device.type == "cuda",
+                    task=TASK_DESCRIPTION,
+                    robot_type=robot.name,
+                )
+
+                # Convert policy output to robot action dict
+                action_values = make_robot_action(action_tensor, dataset.features)
+
+                # Process and send action to robot (EE -> joints via IK)
+                robot_action_to_send = robot_ee_to_joints_processor((action_values, obs))
+                robot.send_action(robot_action_to_send)
+
+                # Write to dataset
+                action_frame = build_dataset_frame(dataset.features, action_values, prefix=ACTION)
+                frame = {**observation_frame, **action_frame, "task": TASK_DESCRIPTION}
+                dataset.add_frame(frame)
+
+                log_rerun_data(observation=obs_processed, action=action_values)
+
+                dt_s = time.perf_counter() - start_loop_t
+                sleep_time_s = control_interval - dt_s
+                if sleep_time_s < 0:
+                    logging.warning(
+                        f"Evaluate loop is running slower ({1 / dt_s:.1f} Hz) than the target FPS ({FPS} Hz)."
+                    )
+                precise_sleep(max(sleep_time_s, 0.0))
+                timestamp = time.perf_counter() - start_episode_t
 
             # Reset the environment if not stopping or re-recording
             if not events["stop_recording"] and (
                 (episode_idx < NUM_EPISODES - 1) or events["rerecord_episode"]
             ):
                 log_say("Reset the environment")
-                record_loop(
-                    robot=robot,
-                    events=events,
-                    fps=FPS,
-                    control_time_s=EPISODE_TIME_SEC,
-                    single_task=TASK_DESCRIPTION,
-                    display_data=True,
-                    teleop_action_processor=make_default_teleop_action_processor(),
-                    robot_action_processor=robot_ee_to_joints_processor,
-                    robot_observation_processor=robot_joints_to_ee_pose_processor,
-                )
+                log_say("Waiting for environment reset, press right arrow key when ready...")
 
             if events["rerecord_episode"]:
                 log_say("Re-record episode")
@@ -190,7 +222,6 @@ def main():
 
             # Save episode
             dataset.save_episode()
-            episode_idx += 1
     finally:
         # Clean up
         log_say("Stop recording")

examples/so100_to_so100_EE/record.py

@@ -62,21 +62,20 @@ def main():
     follower = SO100Follower(follower_config)
     leader = SO100Leader(leader_config)
 
-    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo:
+    #   https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
     follower_kinematics_solver = RobotKinematics(
         urdf_path="./SO101/so101_new_calib.urdf",
         target_frame_name="gripper_frame_link",
         joint_names=list(follower.bus.motors.keys()),
     )
-
-    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
     leader_kinematics_solver = RobotKinematics(
         urdf_path="./SO101/so101_new_calib.urdf",
         target_frame_name="gripper_frame_link",
         joint_names=list(leader.bus.motors.keys()),
     )
 
-    # Build pipeline to convert follower joints to EE observation
+    # Build pipeline to convert follower joints to EE observation.
     follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservation](
         steps=[
             ForwardKinematicsJointsToEE(
@@ -87,7 +86,7 @@ def main():
         to_output=transition_to_observation,
     )
 
-    # Build pipeline to convert leader joints to EE action
+    # Build pipeline to convert leader joints to EE action.
     leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
         steps=[
             ForwardKinematicsJointsToEE(
@@ -98,9 +97,9 @@ def main():
         to_output=transition_to_robot_action,
     )
 
-    # Build pipeline to convert EE action to follower joints
+    # Build pipeline to convert EE action to follower joints (with safety bounds).
     ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
-        [
+        steps=[
             EEBoundsAndSafety(
                 end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
                 max_ee_step_m=0.10,
@@ -115,13 +114,12 @@ def main():
         to_output=transition_to_robot_action,
     )
 
-    # Create the dataset
+    # Create the dataset, deriving features from the pipelines so the on-disk schema
+    # matches exactly what the pipelines produce at runtime.
     dataset = LeRobotDataset.create(
         repo_id=HF_REPO_ID,
         fps=FPS,
         features=combine_feature_dicts(
-            # Run the feature contract of the pipelines
-            # This tells you how the features would look like after the pipeline steps
             aggregate_pipeline_dataset_features(
                 pipeline=leader_joints_to_ee,
                 initial_features=create_initial_features(action=leader.action_features),
@@ -144,7 +142,7 @@ def main():
 
     # Initialize the keyboard listener and rerun visualization
     listener, events = init_keyboard_listener()
-    init_rerun(session_name="recording_phone")
+    init_rerun(session_name="recording_so100_ee")
 
     try:
         if not leader.is_connected or not follower.is_connected:
@@ -160,14 +158,14 @@ def main():
                 robot=follower,
                 events=events,
                 fps=FPS,
+                teleop_action_processor=leader_joints_to_ee,
+                robot_action_processor=ee_to_follower_joints,
+                robot_observation_processor=follower_joints_to_ee,
                 teleop=leader,
                 dataset=dataset,
                 control_time_s=EPISODE_TIME_SEC,
                 single_task=TASK_DESCRIPTION,
                 display_data=True,
-                teleop_action_processor=leader_joints_to_ee,
-                robot_action_processor=ee_to_follower_joints,
-                robot_observation_processor=follower_joints_to_ee,
             )
 
             # Reset the environment if not stopping or re-recording
@@ -179,13 +177,13 @@ def main():
                     robot=follower,
                     events=events,
                     fps=FPS,
+                    teleop_action_processor=leader_joints_to_ee,
+                    robot_action_processor=ee_to_follower_joints,
+                    robot_observation_processor=follower_joints_to_ee,
                     teleop=leader,
                     control_time_s=RESET_TIME_SEC,
                     single_task=TASK_DESCRIPTION,
                     display_data=True,
-                    teleop_action_processor=leader_joints_to_ee,
-                    robot_action_processor=ee_to_follower_joints,
-                    robot_observation_processor=follower_joints_to_ee,
                 )
 
             if events["rerecord_episode"]:

examples/so100_to_so100_EE/rollout.py

@@ -0,0 +1,134 @@
+# !/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Run a trained EE-space policy on SO100 without recording (base rollout).
+
+Uses the rollout engine's :class:`BaseStrategy` (autonomous execution,
+no dataset) with :class:`SyncInferenceConfig` (inline policy call per
+control tick).  The custom observation/action processors convert between
+joint space (robot hardware) and end-effector space (policy I/O) via
+forward/inverse kinematics.
+"""
+
+from lerobot.cameras.opencv import OpenCVCameraConfig
+from lerobot.configs import PreTrainedConfig
+from lerobot.model.kinematics import RobotKinematics
+from lerobot.processor import (
+    RobotProcessorPipeline,
+    observation_to_transition,
+    robot_action_observation_to_transition,
+    transition_to_observation,
+    transition_to_robot_action,
+)
+from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
+from lerobot.robots.so_follower.robot_kinematic_processor import (
+    ForwardKinematicsJointsToEE,
+    InverseKinematicsEEToJoints,
+)
+from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
+from lerobot.rollout.inference import SyncInferenceConfig
+from lerobot.rollout.strategies import BaseStrategy
+from lerobot.types import RobotAction, RobotObservation
+from lerobot.utils.process import ProcessSignalHandler
+from lerobot.utils.utils import init_logging
+
+FPS = 30
+DURATION_SEC = 60
+TASK_DESCRIPTION = "My task description"
+HF_MODEL_ID = "<hf_username>/<model_repo_id>"
+
+
+def main():
+    init_logging()
+
+    # Robot configuration — the rollout engine will connect it inside build_rollout_context.
+    camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
+    robot_config = SO100FollowerConfig(
+        port="/dev/tty.usbmodem5A460814411",
+        id="my_awesome_follower_arm",
+        cameras=camera_config,
+        use_degrees=True,
+    )
+
+    # Kinematic solver: we need the motor-name list, so peek at the robot once.
+    # (The rollout engine owns the connected instance; we only use this for introspection.)
+    temp_robot = SO100Follower(robot_config)
+    motor_names = list(temp_robot.bus.motors.keys())
+
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo:
+    #   https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    kinematics_solver = RobotKinematics(
+        urdf_path="./SO101/so101_new_calib.urdf",
+        target_frame_name="gripper_frame_link",
+        joint_names=motor_names,
+    )
+
+    # Joint-space observation → EE-space observation (consumed by the policy).
+    robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
+        steps=[ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=motor_names)],
+        to_transition=observation_to_transition,
+        to_output=transition_to_observation,
+    )
+
+    # EE-space action (produced by the policy) → joint-space action (sent to robot).
+    robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
+        steps=[
+            InverseKinematicsEEToJoints(
+                kinematics=kinematics_solver,
+                motor_names=motor_names,
+                initial_guess_current_joints=True,
+            ),
+        ],
+        to_transition=robot_action_observation_to_transition,
+        to_output=transition_to_robot_action,
+    )
+
+    # Policy config (full model is loaded inside build_rollout_context).
+    policy_config = PreTrainedConfig.from_pretrained(HF_MODEL_ID)
+    policy_config.pretrained_path = HF_MODEL_ID
+
+    cfg = RolloutConfig(
+        robot=robot_config,
+        policy=policy_config,
+        strategy=BaseStrategyConfig(),
+        inference=SyncInferenceConfig(),
+        fps=FPS,
+        duration=DURATION_SEC,
+        task=TASK_DESCRIPTION,
+    )
+
+    signal_handler = ProcessSignalHandler(use_threads=True)
+
+    # Pass the EE kinematic processors via kwargs; the defaults (identity) would
+    # otherwise skip the joint↔EE conversion and the policy would receive the
+    # wrong observation/action space.
+    ctx = build_rollout_context(
+        cfg,
+        signal_handler.shutdown_event,
+        robot_action_processor=robot_ee_to_joints_processor,
+        robot_observation_processor=robot_joints_to_ee_pose_processor,
+    )
+
+    strategy = BaseStrategy(cfg.strategy)
+    try:
+        strategy.setup(ctx)
+        strategy.run(ctx)
+    finally:
+        strategy.teardown(ctx)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -95,7 +95,7 @@ dependencies = [
 
 # ── Feature-scoped extras ──────────────────────────────────
 dataset = [
-    "datasets>=4.7.0,<5.0.0",
+    "datasets>=4.0.0,<5.0.0",
     "pandas>=2.0.0,<3.0.0", # NOTE: Transitive dependency of datasets
     "pyarrow>=21.0.0,<30.0.0", # NOTE: Transitive dependency of datasets
     "lerobot[av-dep]",
@@ -200,23 +200,6 @@ hilserl = ["lerobot[transformers-dep]", "gym-hil>=0.1.13,<0.2.0", "lerobot[grpci
 async = ["lerobot[grpcio-dep]", "lerobot[matplotlib-dep]"]
 peft = ["lerobot[transformers-dep]", "lerobot[peft-dep]"]
 
-# Annotation pipeline (lerobot-annotate). datatrove is mandatory; vllm is
-# the preferred backend on Linux, with a transformers fallback elsewhere.
-annotations = [
-    "lerobot[dataset]",
-    "lerobot[transformers-dep]",
-    "datatrove>=0.4.0,<2.0.0",
-    "vllm>=0.6.0,<1.0.0; sys_platform == 'linux'",
-]
-
-# Tool implementations under src/lerobot/tools/. Each tool's dependencies
-# are isolated so adding a new tool doesn't bloat the base install.
-# Currently only `say` (Kyutai pocket-tts; CPU-only, ~100M params).
-tools = [
-    "pocket-tts>=0.1.0,<1.0.0",
-    "scipy>=1.11.0,<2.0.0",  # SayTool.output_dir uses scipy.io.wavfile
-]
-
 # Development
 dev = ["pre-commit>=3.7.0,<5.0.0", "debugpy>=1.8.1,<1.9.0", "lerobot[grpcio-dep]", "grpcio-tools==1.73.1", "mypy>=1.19.1", "ruff>=0.14.1", "lerobot[notebook]"]
 notebook = ["jupyter>=1.0.0,<2.0.0", "ipykernel>=6.0.0,<7.0.0"]
@@ -229,20 +212,6 @@ aloha = ["lerobot[dataset]", "gym-aloha>=0.1.2,<0.2.0", "lerobot[scipy-dep]"]
 pusht = ["lerobot[dataset]", "gym-pusht>=0.1.5,<0.2.0", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
 libero = ["lerobot[dataset]", "lerobot[transformers-dep]", "hf-libero>=0.1.3,<0.2.0; sys_platform == 'linux'", "lerobot[scipy-dep]"]
 metaworld = ["lerobot[dataset]", "metaworld==3.0.0", "lerobot[scipy-dep]"]
-# NOTE: vlabench is NOT exposed as a `lerobot` extra. Its only distribution
-# is the OpenMOSS/VLABench GitHub repo (package name `VLABench`, no PyPI
-# release), so any `vlabench>=X` pip spec is unresolvable. Install it
-# manually alongside MuJoCo / dm-control — see docs/source/vlabench.mdx
-# for the recipe.
-# NOTE: robomme is NOT a pyproject extra — mani-skill hard-pins numpy<2
-# which conflicts with lerobot's numpy>=2 base pin, so the two trees can't
-# resolve into a single env. Install it only in the RoboMME Docker image
-# via `uv pip install --override` (see docker/Dockerfile.benchmark.robomme).
-# NOTE: robocasa is NOT exposed as a `lerobot` extra. Its setup.py pins
-# `lerobot==0.3.3` in install_requires, which cyclically shadows our own
-# workspace `lerobot` and makes the graph unsolvable under any resolver
-# (uv, pip). Install it manually alongside robosuite — see
-# docs/source/robocasa.mdx for the recipe.
 
 # All
 all = [
@@ -306,12 +275,11 @@ lerobot-find-joint-limits="lerobot.scripts.lerobot_find_joint_limits:main"
 lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"
 lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main"
 lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main"
-lerobot-annotate="lerobot.scripts.lerobot_annotate:main"
-lerobot-smolvla2-runtime="lerobot.scripts.lerobot_smolvla2_runtime:main"
+lerobot-rollout="lerobot.scripts.lerobot_rollout:main"
 
 # ---------------- Tool Configurations ----------------
 [tool.setuptools.package-data]
-lerobot = ["envs/*.json", "annotations/steerable_pipeline/prompts/*.txt"]
+lerobot = ["envs/*.json"]
 
 [tool.setuptools.packages.find]
 where = ["src"]

scripts/ci/extract_task_descriptions.py

@@ -31,23 +31,9 @@ from __future__ import annotations
 
 import argparse
 import json
-import re
 import sys
 from pathlib import Path
 
-# LIBERO-plus derives task.language by space-joining the perturbation-variant
-# filename (grab_language_from_filename in libero/libero/benchmark/__init__.py),
-# so non-_language_ variants inherit a trailing metadata blob like
-# "view 0 0 100 0 0 initstate 0 noise 45" or "add 16". Strip those tokens so
-# the description matches the base instruction used in the training dataset.
-_LIBERO_PERTURBATION_TAIL_RE = re.compile(
-    r"(?:\s(?:view|initstate|noise|add|tb|table|light|level)(?:\s\d+)+)+$"
-)
-
-
-def _strip_libero_perturbation_tail(instruction: str) -> str:
-    return _LIBERO_PERTURBATION_TAIL_RE.sub("", instruction).strip()
-
 
 def _libero_descriptions(task_suite: str) -> dict[str, str]:
     from libero.libero import benchmark  # type: ignore[import-untyped]
@@ -61,10 +47,7 @@ def _libero_descriptions(task_suite: str) -> dict[str, str]:
         )
         return {}
     suite = suite_dict[task_suite]()
-    return {
-        f"{task_suite}_{i}": _strip_libero_perturbation_tail(suite.get_task(i).language)
-        for i in range(suite.n_tasks)
-    }
+    return {f"{task_suite}_{i}": suite.get_task(i).language for i in range(suite.n_tasks)}
 
 
 def _metaworld_descriptions(task_name: str) -> dict[str, str]:
@@ -74,120 +57,19 @@ def _metaworld_descriptions(task_name: str) -> dict[str, str]:
     return {f"{task_name}_0": label}
 
 
-def _robotwin_descriptions(task_names: str) -> dict[str, str]:
-    """Return descriptions for each requested RoboTwin task. Reads
-    `description/task_instruction/<task>.json` from the RoboTwin clone
-    (cwd is /opt/robotwin in CI). Falls back to the task name if missing."""
-    out: dict[str, str] = {}
-    root = Path("description/task_instruction")
-    for name in (t.strip() for t in task_names.split(",") if t.strip()):
-        desc_file = root / f"{name}.json"
-        desc = name.replace("_", " ")
-        if desc_file.is_file():
-            data = json.loads(desc_file.read_text())
-            full = data.get("full_description") or desc
-            # Strip the schema placeholders ({A}, {a}) — keep the sentence readable.
-            desc = full.replace("<", "").replace(">", "")
-        out[f"{name}_0"] = desc
-    return out
-
-
-def _robocasa_descriptions(task_spec: str) -> dict[str, str]:
-    """For each task in the comma-separated list, emit a cleaned-name label.
-
-    RoboCasa episodes carry their language instruction in the env's
-    `ep_meta['lang']`, populated per reset. Pulling it requires spinning
-    up the full kitchen env per task (~seconds each); we use the task
-    name as the key here and let the eval's episode info carry the
-    actual instruction.
-    """
-    out: dict[str, str] = {}
-    for task in (t.strip() for t in task_spec.split(",") if t.strip()):
-        # Split CamelCase into words: "CloseFridge" → "close fridge".
-        label = "".join(f" {c.lower()}" if c.isupper() else c for c in task).strip()
-        out[f"{task}_0"] = label or task
-    return out
-
-
-_ROBOMME_DESCRIPTIONS = {
-    "BinFill": "Fill the target bin with the correct number of cubes",
-    "PickXtimes": "Pick the indicated cube the specified number of times",
-    "SwingXtimes": "Swing the object the specified number of times",
-    "StopCube": "Grasp and stop the moving cube",
-    "VideoUnmask": "Pick the cube shown in the reference video",
-    "VideoUnmaskSwap": "Pick the cube matching the reference video after a swap",
-    "ButtonUnmask": "Press the button indicated by the reference",
-    "ButtonUnmaskSwap": "Press the correct button after objects are swapped",
-    "PickHighlight": "Pick the highlighted cube",
-    "VideoRepick": "Repick the cube shown in the reference video",
-    "VideoPlaceButton": "Place the cube on the button shown in the video",
-    "VideoPlaceOrder": "Place cubes in the order shown in the video",
-    "MoveCube": "Move the cube to the target location",
-    "InsertPeg": "Insert the peg into the target hole",
-    "PatternLock": "Unlock the pattern by pressing buttons in sequence",
-    "RouteStick": "Route the stick through the required waypoints",
-}
-
-
-def _robomme_descriptions(task_names: str, task_ids: list[int] | None = None) -> dict[str, str]:
-    """Return descriptions for each requested RoboMME task. Keys match the
-    video filename pattern `<task>_<task_id>` used by the eval script."""
-    if task_ids is None:
-        task_ids = [0]
-    out: dict[str, str] = {}
-    for name in (t.strip() for t in task_names.split(",") if t.strip()):
-        desc = _ROBOMME_DESCRIPTIONS.get(name, name)
-        for tid in task_ids:
-            out[f"{name}_{tid}"] = desc
-    return out
-
-
-def _vlabench_descriptions(task_spec: str) -> dict[str, str]:
-    """For each task in the comma-separated list, emit a cleaned-name label.
-
-    VLABench tasks carry language instructions on their dm_control task
-    object, but pulling them requires loading the full env per task
-    (~seconds each). The CI smoke-eval already captures the instruction
-    inside its episode info; this mapping is just enough to key
-    `metrics.json` by `<task>_0`.
-    """
-    out: dict[str, str] = {}
-    for task in (t.strip() for t in task_spec.split(",") if t.strip()):
-        out[f"{task}_0"] = task.replace("_", " ").strip()
-    return out
-
-
 def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--env", required=True, help="Environment family (libero, metaworld, ...)")
     parser.add_argument("--task", required=True, help="Task/suite name (e.g. libero_spatial)")
-    parser.add_argument(
-        "--task-ids",
-        type=str,
-        default=None,
-        help="Comma-separated task IDs (e.g. '0,1,2'). Default: [0]",
-    )
     parser.add_argument("--output", required=True, help="Path to write task_descriptions.json")
     args = parser.parse_args()
 
-    task_ids: list[int] | None = None
-    if args.task_ids:
-        task_ids = [int(x.strip()) for x in args.task_ids.split(",")]
-
     descriptions: dict[str, str] = {}
     try:
-        if args.env == ("libero", "libero_plus"):
+        if args.env == "libero":
             descriptions = _libero_descriptions(args.task)
         elif args.env == "metaworld":
             descriptions = _metaworld_descriptions(args.task)
-        elif args.env == "robotwin":
-            descriptions = _robotwin_descriptions(args.task)
-        elif args.env == "robocasa":
-            descriptions = _robocasa_descriptions(args.task)
-        elif args.env == "robomme":
-            descriptions = _robomme_descriptions(args.task, task_ids=task_ids)
-        elif args.env == "vlabench":
-            descriptions = _vlabench_descriptions(args.task)
         else:
             print(
                 f"[extract_task_descriptions] No description extractor for env '{args.env}'.",

src/lerobot/annotations/__init__.py

@@ -1,15 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

src/lerobot/annotations/steerable_pipeline/__init__.py

@@ -1,36 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Steerable annotation pipeline producing ``language_persistent`` and
-``language_events`` columns for LeRobot datasets.
-
-The pipeline is decomposed into three independently runnable modules whose
-outputs are staged per-episode before a final parquet rewrite:
-
-- :mod:`.modules.plan_subtasks_memory` (Module 1) — persistent styles
-- :mod:`.modules.interjections_and_speech` (Module 2) — event styles + speech
-- :mod:`.modules.general_vqa` (Module 3) — event-style VQA pairs
-"""
-
-from .config import AnnotationPipelineConfig
-from .validator import StagingValidator, ValidationReport
-from .writer import LanguageColumnsWriter
-
-__all__ = [
-    "AnnotationPipelineConfig",
-    "LanguageColumnsWriter",
-    "StagingValidator",
-    "ValidationReport",
-]

src/lerobot/annotations/steerable_pipeline/config.py

@@ -1,260 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any
-
-
-@dataclass
-class Module1Config:
-    """Module 1 hyperparameters: plan + subtasks + memory + task augmentation.
-
-    Subtask decomposition sees the **whole episode** as one Qwen-VL video
-    block — no keyframe stride or count: the model handles temporal pooling
-    itself and decides where to cut. ``max_video_frames`` only caps the
-    number of frames packed into the video block (a model-capacity bound,
-    not an annotation-logic knob).
-    """
-
-    enabled: bool = True
-    n_task_rephrasings: int = 10
-    """Number of task rephrasings to generate at ``t=0`` as ``task_aug``
-    persistent rows (PR 1 ``CORE_STYLES``). The renderer's ``${task}``
-    binding rotates among them deterministically per ``sample_idx``,
-    realizing Xiao 2022 / CAST-style task-prompt diversity without
-    touching ``meta/tasks.parquet``. Set to 0 to disable."""
-    derive_task_from_video: str = "if_short"
-    """When to bypass the user-provided ``record.episode_task`` and
-    derive a fresh task description from the episode video alone:
-
-    - ``off``       never; always use the canonical task as the basis.
-    - ``if_short``  derive when the canonical task is empty, has fewer
-                    than ``derive_task_min_words`` words, or matches a
-                    placeholder string (``debug``, ``unnamed``, ``tbd``,
-                    ...). Default — fixes noisy / placeholder tasks
-                    without forcing derivation everywhere.
-    - ``always``    ignore the canonical task entirely; always derive
-                    from the video. Useful when the dataset's task
-                    labels are uniformly bad.
-
-    The video-derived task replaces the canonical task as the basis for
-    subtask decomposition, plan, memory, AND the ``task_aug`` rephrasings,
-    so every downstream annotation is grounded in what's actually visible.
-    ``meta/tasks.parquet`` is NOT modified — the Module-1-derived task
-    only lives in ``language_persistent`` rows."""
-    derive_task_min_words: int = 3
-    """Word-count threshold for ``derive_task_from_video=if_short``."""
-    frames_per_second: float = 1.0
-    """Sample one image-frame per ``1/fps`` seconds across the episode for
-    Module 1's subtask-decomposition prompt. ``1.0`` = 1 fps. Capped by
-    ``max_video_frames`` to avoid blowing up the request payload."""
-    max_video_frames: int = 128
-    """Hard cap on the number of frames Module 1 sends. With ``fps=1`` and
-    a 30 s episode this yields 30 frames. Bumped from 32 since each frame
-    is small (~30-100 KB PNG when base64'd)."""
-    min_subtask_seconds: float = 1.5
-    plan_max_steps: int = 8
-    use_video_url: bool = False
-    """When True (and backend supports it, e.g. ``openai``), Module 1
-    sends a ``video_url`` content block pointing at the episode's mp4
-    file instead of pre-decoded frames. Lets the server sample frames at
-    its own ``fps`` — no in-process conv3d cost. The video file is
-    extracted as a per-episode subclip to ``staging/.video_clips/`` so
-    the model sees only this episode's frames."""
-    use_video_url_fps: float = 1.0
-    """Frame-rate hint to send to the server (mm_processor_kwargs.fps).
-    Only used when ``use_video_url=True``. ``1.0`` = sample 1 frame per
-    second, which is plenty for subtask-boundary detection on most
-    manipulation episodes."""
-
-
-@dataclass
-class Module2Config:
-    """Module 2 hyperparameters: interjections + paired speech."""
-
-    enabled: bool = True
-    max_interjections_per_episode: int = 3
-    """Number of mid-episode interjections to generate per episode. Each
-    creates a paired ``(interjection, speech)`` event row plus triggers a
-    ``plan`` refresh at the same timestamp via Module 1. Bumped from the
-    original ``1`` after qwen36moe-10 showed plan/interjection coverage
-    was too sparse for Hi Robot-style training."""
-    interjection_min_t: float = 2.0
-    interjection_window_seconds: float = 2.0
-    """How many seconds of video to attach to the interjection prompt as
-    visual context. Without this the VLM only sees a single frozen frame
-    and writes generic interjections that aren't grounded in the actual
-    motion happening at the chosen timestamp."""
-    interjection_window_frames: int = 4
-    """How many frames to sample over ``interjection_window_seconds``.
-    Default 4 ⇒ ~0.5 fps over the leading 2 seconds — enough for the
-    model to read the ongoing motion, cheap enough to keep prompt size
-    bounded for the 32k context."""
-
-
-@dataclass
-class Module3Config:
-    """Module 3 hyperparameters: general VQA."""
-
-    enabled: bool = True
-    vqa_emission_hz: float = 1.0
-    K: int = 3
-    question_types: tuple[str, ...] = ("bbox", "keypoint", "count", "attribute", "spatial")
-
-
-@dataclass
-class VlmConfig:
-    """Shared Qwen-VL client configuration."""
-
-    backend: str = "openai"
-    """One of ``vllm``, ``transformers``, ``openai``, or ``stub`` (tests only).
-
-    Default ``openai`` talks to a local OpenAI-compatible server (vllm /
-    transformers) which the CLI auto-spawns when ``auto_serve=True``."""
-    model_id: str = "Qwen/Qwen2.5-VL-7B-Instruct"
-    api_base: str = "http://localhost:8000/v1"
-    """Base URL for the ``openai`` backend."""
-    api_key: str = "EMPTY"
-    """API key for the ``openai`` backend; ``EMPTY`` works for local servers."""
-    auto_serve: bool = True
-    """When True with ``backend=openai``, the CLI probes ``api_base``
-    first; if no server answers, it spawns one (default:
-    ``transformers serve``), waits for it to be ready, runs the
-    pipeline, and tears it down on exit. Default ``True`` so a single
-    ``lerobot-annotate`` call can drive the whole flow. Set to ``False``
-    if you want to fail fast when no server is reachable (e.g. you're
-    pointing at a remote endpoint that should already be up)."""
-    serve_port: int = 8000
-    """Port the auto-spawned server binds to. Sets ``api_base`` automatically."""
-    serve_command: str | None = None
-    """Override the auto-serve command (full shell command). When ``None``,
-    we run ``transformers serve <model_id> --port <serve_port> --continuous-batching``.
-
-    When ``parallel_servers > 1``, the literal ``{port}`` placeholder in
-    this command (if present) is substituted per-replica."""
-    parallel_servers: int = 1
-    """When >1, spawn this many independent inference servers (each pinned
-    to a GPU via ``CUDA_VISIBLE_DEVICES`` and listening on
-    ``serve_port + i``) and round-robin client requests across them.
-    Useful when DP/TP NCCL setup is broken on the node — single-GPU
-    replicas don't need cross-GPU communication. When
-    ``parallel_servers > num_gpus``, replicas are round-robin-assigned
-    to GPUs (e.g. 4 replicas on 2 GPUs → 0,1,0,1)."""
-    num_gpus: int = 0
-    """How many physical GPUs are available for round-robin replica
-    placement. ``0`` means ``parallel_servers`` (one GPU per replica,
-    backward-compatible default). Set this to ``2`` with
-    ``parallel_servers=4`` to pack 2 replicas per GPU."""
-    client_concurrency: int = 16
-    """Maximum number of in-flight chat requests the client issues in
-    parallel. vllm batches them internally for free, so bumping this
-    typically gives big throughput wins on a single TP=1 server. Set to
-    ``1`` for strict serial calls."""
-    serve_ready_timeout_s: float = 600.0
-    """Max seconds to wait for the server to start serving requests."""
-    max_new_tokens: int = 512
-    temperature: float = 0.2
-    json_mode: bool = True
-    batch_size: int = 4
-    tensor_parallel_size: int = 1
-    gpu_memory_utilization: float = 0.9
-    """Fraction of GPU memory vllm allocates for weights + KV cache.
-    Lower (e.g. 0.7) when the vision encoder needs cuDNN workspace, or to
-    avoid CUDNN_STATUS_NOT_INITIALIZED on tight VRAM (30B BF16 on 80 GB)."""
-    max_model_len: int | None = None
-    """Cap context length. ``None`` keeps the model's default; on H100 80 GB
-    a 30B BF16 model often needs ``max_model_len=8192`` or smaller to leave
-    room for KV cache."""
-    trust_remote_code: bool = False
-    """Pass ``trust_remote_code`` to HF auto-classes. Default ``False`` —
-    only enable for models that actually ship custom code in their repo
-    (rare for first-class VL releases). On Qwen3-VL it triggers an
-    std::bad_alloc post-load even though the official transformers class
-    is sufficient, so leaving this off is safest."""
-    camera_key: str | None = None
-    """Override the camera stream used for keyframe attachment. ``None`` picks
-    the first ``observation.images.*`` key the dataset declares."""
-    chat_template_kwargs: dict[str, Any] | None = None
-    """Forwarded as ``extra_body.chat_template_kwargs`` on every chat call.
-    Use this to pass model-specific template flags such as
-    ``{"enable_thinking": false}`` for Qwen3.5/Qwen3.6 to suppress the
-    reasoning preamble that otherwise eats the entire ``max_new_tokens``
-    budget before any JSON is emitted."""
-
-
-@dataclass
-class ExecutorConfig:
-    """Executor selection and SLURM hyperparameters."""
-
-    auto_threshold: int = 32
-    force_local: bool = False
-    slurm_partition: str | None = None
-    slurm_gpus: int = 1
-    slurm_time: str = "06:00:00"
-    workers: int = 1
-    episode_parallelism: int = 16
-    """Number of episodes processed concurrently within each module phase.
-    Each in-flight episode sends 3–5 dependent VLM calls; bumping this is
-    how you actually saturate ``parallel_servers`` and ``client_concurrency``
-    — without it, the executor loops one episode at a time and the
-    inference servers sit ~90% idle. Set to ``1`` for strict serial
-    execution."""
-
-
-@dataclass
-class AnnotationPipelineConfig:
-    """Top-level config for ``lerobot-annotate``.
-
-    Mirrors the structure of :class:`lerobot.configs.train.TrainPipelineConfig`:
-    a draccus-parsed dataclass that contains nested per-module sub-configs and
-    leaves the dataset, executor, and VLM choices independently knobbable.
-
-    Output is always in-place: the writer rewrites ``data/chunk-*/file-*.parquet``
-    in place. Multiple revisions of the same dataset live in separate copies.
-    """
-
-    repo_id: str | None = None
-    root: Path | None = None
-
-    staging_dir: Path | None = None
-    """If unset, defaults to ``<root>/.annotate_staging/``."""
-
-    seed: int = 1729
-
-    module_1: Module1Config = field(default_factory=Module1Config)
-    module_2: Module2Config = field(default_factory=Module2Config)
-    module_3: Module3Config = field(default_factory=Module3Config)
-
-    vlm: VlmConfig = field(default_factory=VlmConfig)
-    executor: ExecutorConfig = field(default_factory=ExecutorConfig)
-
-    skip_validation: bool = False
-    only_episodes: tuple[int, ...] | None = None
-
-    push_to_hub: str | None = None
-    """If set, after the pipeline completes, upload the annotated dataset
-    root to the Hugging Face Hub as a dataset repo with this id (e.g.
-    ``pepijn/super_poulain_steerable``). Creates the repo if missing."""
-    push_private: bool = False
-    """When ``push_to_hub`` is set, create the repo as private."""
-    push_commit_message: str | None = None
-    """Override the commit message used for the hub upload."""
-
-    def resolved_staging_dir(self, root: Path) -> Path:
-        return self.staging_dir if self.staging_dir is not None else root / ".annotate_staging"

src/lerobot/annotations/steerable_pipeline/executor.py

@@ -1,263 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Executor selection: local vs SLURM via datatrove.
-
-The executor plans **four phases** with the dependency order from the plan:
-
-    phase 1: Module 1 (plan + subtasks + memory)
-    phase 2: Module 2 (interjections + speech)
-    phase 3: Module 1 plan-update pass — re-runs plan emission at every
-             interjection timestamp produced by phase 2
-    phase 4: Module 3 (VQA)
-    phase 5: validator
-    phase 6: writer
-
-Phase 3 is why ``executor.py`` documents the dependency: Module 1 must be
-re-entered after Module 2 to refresh ``plan`` rows at interjection times.
-"""
-
-from __future__ import annotations
-
-import logging
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-from .config import AnnotationPipelineConfig, ExecutorConfig
-from .reader import EpisodeRecord, iter_episodes
-from .staging import EpisodeStaging
-from .validator import StagingValidator
-from .writer import LanguageColumnsWriter
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class PhaseResult:
-    """Summary of one pipeline phase across all episodes."""
-
-    name: str
-    episodes_processed: int
-    episodes_skipped: int
-
-
-@dataclass
-class PipelineRunSummary:
-    """Aggregated result returned by :meth:`Executor.run`."""
-
-    phases: list[PhaseResult]
-    written_paths: list[Path]
-    validation_report: Any  # ValidationReport, kept Any to avoid import cycle
-
-
-def select_executor_class(num_episodes: int, config: ExecutorConfig) -> str:
-    """Return ``"local"`` or ``"slurm"`` based on the threshold.
-
-    The plan's "executor selection threshold" lives in
-    :class:`ExecutorConfig.auto_threshold`. ``force_local`` always wins.
-    """
-    if config.force_local:
-        return "local"
-    return "local" if num_episodes <= config.auto_threshold else "slurm"
-
-
-@dataclass
-class Executor:
-    """Run all four phases over a dataset root.
-
-    The executor is intentionally framework-agnostic: by default it runs the
-    phases inline (suitable for tests, small datasets, and the CLI's
-    ``--force-local`` mode). It will optionally hand off to datatrove's
-    :class:`LocalPipelineExecutor` or :class:`SlurmPipelineExecutor` when those
-    are installed and the dataset is large enough to benefit from them.
-
-    Tests construct the executor directly with stub modules.
-    """
-
-    config: AnnotationPipelineConfig
-    module_1: Any  # PlanSubtasksMemoryModule
-    module_2: Any  # InterjectionsAndSpeechModule
-    module_3: Any  # GeneralVqaModule
-    writer: LanguageColumnsWriter
-    validator: StagingValidator
-
-    def run(self, root: Path) -> PipelineRunSummary:
-        records = list(iter_episodes(root, only_episodes=self.config.only_episodes))
-        n = len(records)
-        if n == 0:
-            raise ValueError(f"No episodes found under {root}/data/")
-
-        executor_kind = select_executor_class(n, self.config.executor)
-        print(f"[annotate] {n} episodes total; executor={executor_kind}", flush=True)
-
-        staging_dir = self.config.resolved_staging_dir(root)
-        staging_dir.mkdir(parents=True, exist_ok=True)
-
-        phases: list[PhaseResult] = []
-
-        # Phase 1: Module 1 (plan + subtasks + memory)
-        phases.append(self._run_module_phase("module_1", records, staging_dir, self.module_1))
-        # Phase 2: Module 2 (interjections + speech). Module 2 reads
-        # Module 1's subtask rows from the same staging tree to ground
-        # the interjection prompt in the correct local subtask.
-        phases.append(self._run_module_phase("module_2", records, staging_dir, self.module_2))
-        # Phase 3: Module 1 plan-update pass at interjection timestamps.
-        phases.append(self._run_plan_update_phase(records, staging_dir))
-        # Phase 4: Module 3 (VQA)
-        phases.append(self._run_module_phase("module_3", records, staging_dir, self.module_3))
-
-        print("[annotate] running validator...", flush=True)
-        report = self.validator.validate(records, staging_dir)
-        if not report.ok and not self.config.skip_validation:
-            raise RuntimeError(f"Staging validation failed: {report.summary()}")
-        print(f"[annotate] validator: {report.summary()}", flush=True)
-
-        print(f"[annotate] writing parquet shards into {root}/data/...", flush=True)
-        written = self.writer.write_all(records, staging_dir, root)
-        print(f"[annotate] wrote {len(written)} shard(s); pipeline complete", flush=True)
-
-        # Persist the tool catalog to meta/info.json so chat-template
-        # consumers (PR 3 SmolVLA2 / Pi0.5 / dataset visualizer) can read
-        # it via ``LeRobotDatasetMetadata.tools`` (PR 1). Idempotent and
-        # additive: anything the user pre-populated is preserved; we only
-        # ensure the canonical ``say`` schema is present.
-        self._ensure_tools_in_info(root)
-
-        return PipelineRunSummary(phases=phases, written_paths=written, validation_report=report)
-
-    def _ensure_tools_in_info(self, root: Path) -> None:
-        """Write ``meta/info.json["tools"]`` if missing the canonical ``say``.
-
-        Reads any user-declared tools already in ``info.json`` and merges
-        the canonical ``SAY_TOOL_SCHEMA`` into the list (deduped by
-        ``function.name``). Writes back to disk only if the list
-        changed.
-        """
-        import json  # noqa: PLC0415
-
-        from lerobot.datasets.language import SAY_TOOL_SCHEMA  # noqa: PLC0415
-
-        info_path = root / "meta" / "info.json"
-        if not info_path.exists():
-            return
-        try:
-            info = json.loads(info_path.read_text())
-        except Exception as exc:  # noqa: BLE001
-            print(f"[annotate] could not read {info_path}: {exc}", flush=True)
-            return
-
-        existing = info.get("tools")
-        if not isinstance(existing, list):
-            existing = []
-        names = {
-            (t.get("function") or {}).get("name")
-            for t in existing
-            if isinstance(t, dict)
-        }
-        merged = list(existing)
-        if SAY_TOOL_SCHEMA["function"]["name"] not in names:
-            merged.append(SAY_TOOL_SCHEMA)
-        if merged != existing:
-            info["tools"] = merged
-            info_path.write_text(json.dumps(info, indent=2))
-            print(
-                f"[annotate] meta/info.json: tools={[t['function']['name'] for t in merged]}",
-                flush=True,
-            )
-
-    def _run_module_phase(
-        self,
-        name: str,
-        records: list[EpisodeRecord],
-        staging_dir: Path,
-        module: Any,
-    ) -> PhaseResult:
-        import time as _time  # noqa: PLC0415
-        from concurrent.futures import ThreadPoolExecutor, as_completed  # noqa: PLC0415
-
-        if not module.enabled:
-            print(f"[annotate] phase={name} skipped (module disabled)", flush=True)
-            return PhaseResult(name=name, episodes_processed=0, episodes_skipped=len(records))
-        n = len(records)
-        parallelism = max(1, min(self.config.executor.episode_parallelism, n))
-        print(
-            f"[annotate] phase={name} starting on {n} episode(s) "
-            f"(parallelism={parallelism})",
-            flush=True,
-        )
-        t0 = _time.time()
-
-        def _do(idx_record: tuple[int, EpisodeRecord]) -> tuple[int, int, float]:
-            i, record = idx_record
-            ep_start = _time.time()
-            staging = EpisodeStaging(staging_dir, record.episode_index)
-            module.run_episode(record, staging)
-            return i, record.episode_index, _time.time() - ep_start
-
-        processed = 0
-        if parallelism == 1:
-            for i, record in enumerate(records, 1):
-                _, ep_idx, elapsed = _do((i, record))
-                processed += 1
-                print(
-                    f"[annotate]   {name} episode {i}/{n} "
-                    f"(idx={ep_idx}) done in {elapsed:.1f}s",
-                    flush=True,
-                )
-        else:
-            with ThreadPoolExecutor(max_workers=parallelism) as pool:
-                futures = [pool.submit(_do, (i, r)) for i, r in enumerate(records, 1)]
-                for fut in as_completed(futures):
-                    i, ep_idx, elapsed = fut.result()
-                    processed += 1
-                    print(
-                        f"[annotate]   {name} episode {processed}/{n} "
-                        f"(idx={ep_idx}, submit_order={i}) done in {elapsed:.1f}s",
-                        flush=True,
-                    )
-        total = _time.time() - t0
-        print(f"[annotate] phase={name} complete: {processed}/{n} in {total:.1f}s", flush=True)
-        return PhaseResult(name=name, episodes_processed=processed, episodes_skipped=0)
-
-    def _run_plan_update_phase(  # noqa: PLR0915
-        self, records: list[EpisodeRecord], staging_dir: Path
-    ) -> PhaseResult:
-        """Re-emit ``plan`` rows at each interjection timestamp from Module 2.
-
-        Module 1 owns the prompt; Module 2 produced the timestamps. This phase
-        therefore calls back into Module 1 with the interjection timestamps so
-        Module 1's existing prompt path is reused.
-        """
-        if not self.module_1.enabled or not self.module_2.enabled:
-            return PhaseResult(
-                name="module_1_plan_update", episodes_processed=0, episodes_skipped=len(records)
-            )
-        processed = 0
-        for record in records:
-            staging = EpisodeStaging(staging_dir, record.episode_index)
-            interjection_rows = [
-                row
-                for row in staging.read("module_2")
-                if row.get("style") == "interjection"
-            ]
-            interjection_times = [float(row["timestamp"]) for row in interjection_rows]
-            interjection_texts = [str(row.get("content") or "") for row in interjection_rows]
-            if interjection_times:
-                self.module_1.run_plan_updates(
-                    record, staging, interjection_times, interjection_texts
-                )
-                processed += 1
-        return PhaseResult(name="module_1_plan_update", episodes_processed=processed, episodes_skipped=0)

src/lerobot/annotations/steerable_pipeline/frames.py

@@ -1,400 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Keyframe extraction for the annotation pipeline.
-
-Modules attach decoded camera frames to their VLM prompts so the model can
-ground subtask decomposition, interjection scenarios, and VQA in actual
-visual content. The pipeline shares one provider across modules and one
-episode at a time, with a small per-episode cache so multiple modules
-querying the same timestamp pay decode cost once.
-"""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any, Protocol
-
-from .reader import EpisodeRecord
-
-
-class FrameProvider(Protocol):
-    """Decodes camera frames at episode-relative timestamps."""
-
-    @property
-    def camera_keys(self) -> list[str]:
-        """All ``observation.images.*`` feature keys this provider can decode."""
-
-    def frames_at(
-        self,
-        record: EpisodeRecord,
-        timestamps: list[float],
-        camera_key: str | None = None,
-    ) -> list[Any]:
-        """Return one PIL.Image per timestamp from ``camera_key`` (or default).
-
-        Empty list if the camera is unavailable. ``camera_key=None`` falls back
-        to the provider's default camera so existing single-camera callers
-        (Module 1, Module 2) keep working unchanged.
-        """
-
-    def video_for_episode(
-        self,
-        record: EpisodeRecord,
-        max_frames: int,
-        camera_key: str | None = None,
-    ) -> list[Any]:
-        """Return up to ``max_frames`` PIL images covering the whole episode.
-
-        Sampling is uniform across the episode duration. The returned list is
-        intended to be passed as one ``{"type":"video", "video":<list>}``
-        block to a Qwen-VL-compatible model that pools temporally itself.
-        Empty list if no camera available.
-        """
-
-
-@dataclass
-class _NullProvider:
-    """No-op provider used when the dataset has no video keys or in tests."""
-
-    @property
-    def camera_keys(self) -> list[str]:
-        return []
-
-    def frames_at(
-        self,
-        record: EpisodeRecord,
-        timestamps: list[float],
-        camera_key: str | None = None,
-    ) -> list[Any]:
-        return []
-
-    def video_for_episode(
-        self,
-        record: EpisodeRecord,
-        max_frames: int,
-        camera_key: str | None = None,
-    ) -> list[Any]:
-        return []
-
-
-def null_provider() -> FrameProvider:
-    return _NullProvider()
-
-
-@dataclass
-class VideoFrameProvider:
-    """Decodes frames from the dataset's ``observation.images.*`` streams.
-
-    By default the *first* camera key is used for Module 1 (subtask
-    decomposition) and Module 2 (interjection scenarios) — those prompts care
-    about *what is happening*, not which angle. Module 3 (VQA) instead
-    iterates over every camera in :attr:`camera_keys` so each frame's
-    grounded answer (bbox/keypoint/...) is tagged with the camera it was
-    grounded against.
-
-    ``camera_key`` overrides the default-camera choice but does not restrict
-    :attr:`camera_keys`. Pass ``camera_key`` explicitly to ``frames_at`` /
-    ``video_for_episode`` to read a non-default stream.
-
-    Caches up to ``cache_size`` decoded frames per process to keep
-    co-timestamped Module 2 + Module 1 plan-update calls cheap.
-    """
-
-    root: Path
-    camera_key: str | None = None
-    tolerance_s: float = 1e-2
-    cache_size: int = 256
-    _meta: Any = field(default=None, init=False, repr=False)
-    _cache: dict = field(default_factory=dict, init=False, repr=False)
-    _camera_keys: list[str] = field(default_factory=list, init=False, repr=False)
-
-    def __post_init__(self) -> None:
-        from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata  # noqa: PLC0415
-
-        self._meta = LeRobotDatasetMetadata(repo_id="local", root=self.root)
-        # ``camera_keys`` covers both image- and video-stored cameras
-        # (``video_keys`` is video-only). Some datasets declare cameras with
-        # ``dtype=image``, which would otherwise look empty here and silently
-        # disable Module 3 even though the videos are there.
-        keys = list(getattr(self._meta, "camera_keys", None) or self._meta.video_keys or [])
-        # Last-resort fallback: if metadata didn't surface anything but the
-        # caller explicitly named a camera (``--vlm.camera_key=...``), trust
-        # them — the key is by definition known to exist on the dataset.
-        if not keys and self.camera_key:
-            keys = [self.camera_key]
-        self._camera_keys = keys
-        if self.camera_key is None:
-            self.camera_key = keys[0] if keys else None
-
-    @property
-    def camera_keys(self) -> list[str]:
-        """All ``observation.images.*`` keys available on this dataset."""
-        return list(self._camera_keys)
-
-    def frames_at(
-        self,
-        record: EpisodeRecord,
-        timestamps: list[float],
-        camera_key: str | None = None,
-    ) -> list[Any]:
-        target = camera_key if camera_key is not None else self.camera_key
-        if not timestamps or target is None:
-            return []
-
-        out: list[Any] = []
-        misses: list[float] = []
-        miss_indices: list[int] = []
-        for i, ts in enumerate(timestamps):
-            key = (record.episode_index, target, round(float(ts), 6))
-            cached = self._cache.get(key)
-            if cached is not None:
-                out.append(cached)
-            else:
-                out.append(None)
-                misses.append(float(ts))
-                miss_indices.append(i)
-
-        if misses:
-            decoded = self._decode(record.episode_index, misses, target)
-            # decoder may return fewer frames than requested when some
-            # timestamps fall outside the video; pair what we have and
-            # leave the rest as None to be filtered below.
-            for i, img in zip(miss_indices, decoded):
-                out[i] = img
-                key = (record.episode_index, target, round(float(timestamps[i]), 6))
-                if len(self._cache) >= self.cache_size:
-                    self._cache.pop(next(iter(self._cache)))
-                self._cache[key] = img
-        # filter out any None left over from decode failures
-        return [img for img in out if img is not None]
-
-    def _decode(
-        self, episode_index: int, timestamps: list[float], camera_key: str
-    ) -> list[Any]:
-        ep = self._meta.episodes[episode_index]
-        from_timestamp = ep[f"videos/{camera_key}/from_timestamp"]
-        shifted = [from_timestamp + ts for ts in timestamps]
-        video_path = self.root / self._meta.get_video_file_path(episode_index, camera_key)
-
-        try:
-            return _decode_pyav_direct(video_path, shifted, self.tolerance_s)
-        except Exception as exc:
-            # Log loudly the first time decoding fails so silent
-            # Module-3-no-op (every prompt skipped because frames_at returned
-            # []) is debuggable from the job log instead of post-hoc parquet
-            # inspection. Subsequent failures stay quiet.
-            if not getattr(self, "_warned_decode_fail", False):
-                import logging  # noqa: PLC0415
-
-                logging.getLogger(__name__).warning(
-                    "VideoFrameProvider._decode failed for episode=%s camera=%s "
-                    "video_path=%s: %s",
-                    episode_index,
-                    camera_key,
-                    video_path,
-                    exc,
-                    exc_info=True,
-                )
-                self._warned_decode_fail = True
-            return []
-
-
-def _decode_pyav_direct(
-    video_path: Any, timestamps: list[float], tolerance_s: float
-) -> list[Any]:
-    """Decode the requested timestamps from ``video_path`` using PyAV directly.
-
-    Bypasses ``lerobot.datasets.video_utils.decode_video_frames`` entirely
-    because its "pyav" path actually goes through
-    ``decode_video_frames_torchvision`` → ``torchvision.io.VideoReader``,
-    which was removed in torchvision >= 0.22 (the vllm/vllm-openai:latest
-    container ships with torchvision 0.25). The annotation pipeline only
-    needs a handful of PIL images per (episode, ts), so we can decode them
-    with PyAV without any torch dependency at all.
-
-    Returns one ``PIL.Image`` per requested timestamp, in the same order.
-    Any timestamp the decoder couldn't reach is silently dropped (mirrors
-    the previous behaviour); callers filter ``None``/missing entries.
-    """
-    import av  # noqa: PLC0415
-    from PIL import Image  # noqa: PLC0415
-
-    if not timestamps:
-        return []
-
-    targets = sorted(set(timestamps))
-    seek_to = max(0.0, min(targets) - max(0.5, tolerance_s))
-
-    container = av.open(str(video_path))
-    try:
-        stream = container.streams.video[0]
-        # PyAV needs the seek target in stream timebase ticks.
-        if stream.time_base is None:
-            seek_pts = 0
-        else:
-            seek_pts = int(seek_to / float(stream.time_base))
-        try:
-            container.seek(seek_pts, any_frame=False, backward=True, stream=stream)
-        except av.AVError:
-            # Some streams reject the explicit seek; fall back to decoding from start.
-            container.seek(0)
-
-        results: dict[float, Any] = {}
-        target_iter = iter(targets)
-        next_target = next(target_iter, None)
-        for frame in container.decode(stream):
-            if next_target is None:
-                break
-            ts = float(frame.pts * frame.time_base) if frame.pts is not None else None
-            if ts is None:
-                continue
-            # Walk past targets we've already overshot — we keep the closest
-            # frame within tolerance.
-            while next_target is not None and ts >= next_target - tolerance_s:
-                if abs(ts - next_target) <= tolerance_s or ts >= next_target:
-                    img = frame.to_image()  # PIL.Image.Image (RGB)
-                    results.setdefault(next_target, img)
-                    next_target = next(target_iter, None)
-                else:
-                    break
-    finally:
-        container.close()
-
-    return [results[ts] for ts in timestamps if ts in results]
-
-    def video_for_episode(
-        self,
-        record: EpisodeRecord,
-        max_frames: int,
-        camera_key: str | None = None,
-    ) -> list[Any]:
-        """Return up to ``max_frames`` images uniformly sampled across the episode.
-
-        The whole episode duration is covered; the model picks subtask
-        boundaries from the temporal pooling it does internally.
-        """
-        target = camera_key if camera_key is not None else self.camera_key
-        if max_frames <= 0 or target is None or not record.frame_timestamps:
-            return []
-        n_frames = min(max_frames, len(record.frame_timestamps))
-        if n_frames == len(record.frame_timestamps):
-            timestamps = list(record.frame_timestamps)
-        else:
-            t0 = record.frame_timestamps[0]
-            t_last = record.frame_timestamps[-1]
-            if t_last <= t0:
-                timestamps = [float(t0)] * n_frames
-            else:
-                step = (t_last - t0) / (n_frames - 1) if n_frames > 1 else 0.0
-                timestamps = [float(t0 + i * step) for i in range(n_frames)]
-        return self.frames_at(record, timestamps, camera_key=target)
-
-
-def make_frame_provider(root: Path, camera_key: str | None = None) -> FrameProvider:
-    """Build a :class:`VideoFrameProvider` if videos are present, else null."""
-    try:
-        provider = VideoFrameProvider(root=root, camera_key=camera_key)
-    except Exception:
-        return null_provider()
-    if provider.camera_key is None:
-        return null_provider()
-    return provider
-
-
-def to_image_blocks(images: list[Any]) -> list[dict[str, Any]]:
-    """Convert PIL images to Qwen-VL-compatible content blocks."""
-    return [{"type": "image", "image": img} for img in images]
-
-
-def to_video_block(images: list[Any]) -> list[dict[str, Any]]:
-    """Wrap a list of PIL images as one Qwen-VL video block.
-
-    Returns ``[]`` when the list is empty, so the caller can splat the result
-    into a content array without a separate emptiness check.
-    """
-    if not images:
-        return []
-    return [{"type": "video", "video": list(images)}]
-
-
-def to_video_url_block(url: str | None, fps: float = 2.0) -> list[dict[str, Any]]:
-    """Wrap a video file URL as one ``video_url`` block.
-
-    Used by the ``openai`` backend (transformers serve / vllm serve /
-    ktransformers serve), where the server handles frame sampling.
-    Returns ``[]`` when ``url`` is ``None`` so the caller can splat.
-    """
-    if not url:
-        return []
-    return [{"type": "video_url", "video_url": {"url": url}, "fps": fps}]
-
-
-def episode_clip_path(
-    record: EpisodeRecord,
-    provider: "VideoFrameProvider",
-    cache_dir: Path,
-) -> Path | None:
-    """Extract the episode's subclip to ``cache_dir/ep_{idx:06d}.mp4``.
-
-    Returns ``None`` if the dataset has no video tracks. Skips re-extract
-    when the cached clip already exists. Re-encodes to H.264
-    (libx264) so the resulting mp4 is decodable by every downstream
-    video processor — stream-copy would inherit the source codec
-    (often AV1 in modern LeRobot datasets), which vllm's libav build
-    cannot decode.
-    """
-    import subprocess  # noqa: PLC0415
-
-    if provider.camera_key is None:
-        return None
-    cache_dir.mkdir(parents=True, exist_ok=True)
-    out_path = cache_dir / f"ep_{record.episode_index:06d}.mp4"
-    if out_path.exists() and out_path.stat().st_size > 0:
-        return out_path
-    ep = provider._meta.episodes[record.episode_index]
-    from_timestamp = float(ep[f"videos/{provider.camera_key}/from_timestamp"])
-    to_timestamp = float(ep[f"videos/{provider.camera_key}/to_timestamp"])
-    src = provider.root / provider._meta.get_video_file_path(
-        record.episode_index, provider.camera_key
-    )
-    cmd = [
-        "ffmpeg",
-        "-y",
-        "-loglevel",
-        "error",
-        "-ss",
-        f"{from_timestamp:.3f}",
-        "-to",
-        f"{to_timestamp:.3f}",
-        "-i",
-        str(src),
-        "-c:v",
-        "libx264",
-        "-preset",
-        "ultrafast",
-        "-crf",
-        "23",
-        "-pix_fmt",
-        "yuv420p",
-        "-an",
-        str(out_path),
-    ]
-    try:
-        subprocess.run(cmd, check=True, timeout=300)
-    except (subprocess.CalledProcessError, subprocess.TimeoutExpired, FileNotFoundError):
-        return None
-    return out_path if out_path.exists() and out_path.stat().st_size > 0 else None

src/lerobot/annotations/steerable_pipeline/modules/__init__.py

@@ -1,25 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from .general_vqa import GeneralVqaModule
-from .interjections_and_speech import InterjectionsAndSpeechModule
-from .plan_subtasks_memory import PlanSubtasksMemoryModule
-
-__all__ = [
-    "GeneralVqaModule",
-    "InterjectionsAndSpeechModule",
-    "PlanSubtasksMemoryModule",
-]

src/lerobot/annotations/steerable_pipeline/modules/general_vqa.py

@@ -1,238 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Module 3: general VQA at a timed cadence.
-
-Anchors ``K`` (question, answer) pairs to ``K`` consecutive frames per
-emission. For datasets with multiple cameras, every emission tick produces
-one ``(vqa, user)`` + ``(vqa, assistant)`` pair *per camera*: each pair is
-generated against that camera's frame and stamped with the matching
-``camera`` field on the emitted rows. The resolver disambiguates via
-``camera=...``; recipes that consume VQA do so through one sub-recipe
-per camera (see ``recipes/pi05_hirobot.yaml``).
-
-Within a single (frame, camera) we still emit at most one ``(vqa, user)``
-and one ``(vqa, assistant)`` row, so the resolver contract stays scalar.
-
-Question types covered (per the plan's Module 3 table): bbox, keypoint,
-count, attribute, spatial. The assistant's ``content`` is a JSON string
-whose schema depends on the question type. Malformed JSON triggers one
-retry inside :meth:`VlmClient.generate_json`.
-"""
-
-from __future__ import annotations
-
-import json
-import random
-from collections.abc import Sequence
-from dataclasses import dataclass, field
-from typing import Any
-
-from ..config import Module3Config
-from ..frames import FrameProvider, null_provider, to_image_blocks
-from ..prompts import load as load_prompt
-from ..reader import EpisodeRecord
-from ..staging import EpisodeStaging
-from ..validator import classify_vqa_answer
-from ..vlm_client import VlmClient
-
-
-def _emission_anchor_indices(frame_timestamps: Sequence[float], hz: float, k: int) -> list[int]:
-    """Return the relative frame indices to anchor VQA emissions to.
-
-    For each emission tick (every ``1/hz`` seconds), we anchor ``k``
-    consecutive frames starting at the tick. Ticks fall on the nearest
-    available source frame timestamp.
-    """
-    if hz <= 0 or k <= 0 or not frame_timestamps:
-        return []
-    t0 = frame_timestamps[0]
-    t_last = frame_timestamps[-1]
-    period = 1.0 / hz
-    indices: list[int] = []
-    t = t0
-    while t <= t_last + 1e-9:
-        # find the index of the nearest frame to t
-        nearest_i = min(range(len(frame_timestamps)), key=lambda i: abs(frame_timestamps[i] - t))
-        for offset in range(k):
-            j = nearest_i + offset
-            if j >= len(frame_timestamps):
-                break
-            if not indices or indices[-1] != j:
-                indices.append(j)
-        t += period
-    # dedupe while preserving order
-    seen: set[int] = set()
-    deduped: list[int] = []
-    for i in indices:
-        if i in seen:
-            continue
-        seen.add(i)
-        deduped.append(i)
-    return deduped
-
-
-@dataclass
-class GeneralVqaModule:
-    """Emit grounded VQA pairs at a timed cadence."""
-
-    vlm: VlmClient
-    config: Module3Config
-    seed: int = 1729
-    frame_provider: FrameProvider = field(default_factory=null_provider)
-
-    @property
-    def enabled(self) -> bool:
-        return self.config.enabled
-
-    def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
-        if not record.frame_timestamps:
-            staging.write("module_3", [])
-            return
-        rng = random.Random(f"{self.seed}:{record.episode_index}:vqa")
-        anchor_idx = _emission_anchor_indices(
-            record.frame_timestamps, self.config.vqa_emission_hz, self.config.K
-        )
-        cameras = self._target_cameras()
-        if not cameras:
-            # No camera available — emit nothing rather than producing
-            # untagged rows that would fail validation. Surface a loud one-
-            # time warning so this is never silently a no-op.
-            if not getattr(self, "_warned_no_camera", False):
-                import logging  # noqa: PLC0415
-
-                logging.getLogger(__name__).warning(
-                    "Module 3 (VQA) found no cameras on the frame provider — "
-                    "every episode will emit zero VQA rows. Check that the "
-                    "dataset declares observation.images.* features in "
-                    "meta/info.json; passing --vlm.camera_key=<key> at the "
-                    "CLI now also seeds the cameras list as a fallback."
-                )
-                self._warned_no_camera = True
-            staging.write("module_3", [])
-            return
-
-        # Build all messages first (one per (frame, camera)), then issue them
-        # as a single batched generate_json call so the client can fan them
-        # out concurrently.
-        per_call: list[tuple[float, str, str, list[dict[str, Any]]]] = []
-        for idx in anchor_idx:
-            ts = float(record.frame_timestamps[idx])
-            qtype = rng.choice(self.config.question_types)
-            for camera in cameras:
-                messages = self._build_messages(record, qtype, ts, camera)
-                # Skip cameras that decoded to zero frames at this ts: no point
-                # asking the VLM to ground a bbox without an image.
-                if not _has_image_block(messages):
-                    continue
-                per_call.append((ts, camera, qtype, messages))
-
-        if not per_call:
-            staging.write("module_3", [])
-            return
-
-        results = self.vlm.generate_json([m for _, _, _, m in per_call])
-
-        rows: list[dict[str, Any]] = []
-        for (ts, camera, _qtype, _messages), result in zip(per_call, results):
-            qa = self._postprocess(result)
-            if qa is None:
-                continue
-            question, answer = qa
-            rows.append(
-                {
-                    "role": "user",
-                    "content": question,
-                    "style": "vqa",
-                    "timestamp": ts,
-                    "camera": camera,
-                    "tool_calls": None,
-                }
-            )
-            rows.append(
-                {
-                    "role": "assistant",
-                    "content": json.dumps(answer, sort_keys=True),
-                    "style": "vqa",
-                    "timestamp": ts,
-                    "camera": camera,
-                    "tool_calls": None,
-                }
-            )
-        staging.write("module_3", rows)
-
-    def _target_cameras(self) -> list[str]:
-        """Return the cameras Module 3 should iterate per emission tick.
-
-        Defaults to every camera the provider exposes. Datasets with no
-        cameras (or test/null providers) yield an empty list, which makes
-        ``run_episode`` a no-op.
-        """
-        return list(getattr(self.frame_provider, "camera_keys", []) or [])
-
-    def _build_messages(
-        self,
-        record: EpisodeRecord,
-        question_type: str,
-        frame_timestamp: float,
-        camera_key: str,
-    ) -> list[dict[str, Any]]:
-        prompt = load_prompt("module_3_vqa").format(
-            episode_task=record.episode_task,
-            question_type=question_type,
-        )
-        images = self.frame_provider.frames_at(
-            record, [frame_timestamp], camera_key=camera_key
-        )
-        content = [*to_image_blocks(images), {"type": "text", "text": prompt}]
-        return [{"role": "user", "content": content}]
-
-    def _postprocess(self, result: Any) -> tuple[str, dict[str, Any]] | None:
-        if not isinstance(result, dict):
-            return None
-        question = result.get("question")
-        answer = result.get("answer")
-        if not isinstance(question, str) or not question.strip():
-            return None
-        if not isinstance(answer, dict):
-            return None
-        # The validator will enforce shape; here we just sanity-check that the
-        # answer matches *some* known shape so we can drop garbage early.
-        if classify_vqa_answer(answer) is None:
-            return None
-        return question.strip(), answer
-
-    def _generate_one(
-        self,
-        record: EpisodeRecord,
-        question_type: str,
-        frame_timestamp: float,
-        camera_key: str,
-    ) -> tuple[str, dict[str, Any]] | None:
-        messages = self._build_messages(record, question_type, frame_timestamp, camera_key)
-        result = self.vlm.generate_json([messages])[0]
-        return self._postprocess(result)
-
-
-def _has_image_block(messages: list[dict[str, Any]]) -> bool:
-    """Return True if any user content block is a populated image block."""
-    for msg in messages:
-        content = msg.get("content")
-        if not isinstance(content, list):
-            continue
-        for block in content:
-            if isinstance(block, dict) and block.get("type") == "image":
-                return True
-    return False

src/lerobot/annotations/steerable_pipeline/modules/interjections_and_speech.py

@@ -1,231 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Module 2: interjections + paired speech (EVENT styles + speech atoms).
-
-Two sub-passes:
-
-1. At ``t=0``, emit ONLY a speech tool-call atom (acknowledgement of the
-   canonical task). No interjection row — the canonical task is already the
-   user utterance from ``meta/tasks.parquet``.
-
-2. For mid-episode interruptions, emit a co-timestamped pair:
-       {role:user, style:interjection, content:<text>}
-       speech atom (role:assistant, style:None, tool_calls=[say(...)])
-   Both rows go in ``language_events`` at the same timestamp.
-
-Module 1's :meth:`run_plan_updates` reuses Module 2's interjection
-timestamps to refresh the ``plan`` row at the same instant.
-"""
-
-from __future__ import annotations
-
-import random
-from collections.abc import Sequence
-from dataclasses import dataclass, field
-from typing import Any
-
-from ..config import Module2Config
-from ..frames import FrameProvider, null_provider, to_image_blocks
-from ..prompts import load as load_prompt
-from ..reader import EpisodeRecord
-from ..staging import EpisodeStaging
-from ..vlm_client import VlmClient
-from ..writer import speech_atom
-
-
-def _snap_to_frame(t: float, frame_timestamps: Sequence[float]) -> float:
-    if not frame_timestamps:
-        return float(t)
-    return float(min(frame_timestamps, key=lambda f: abs(f - t)))
-
-
-@dataclass
-class InterjectionsAndSpeechModule:
-    """Generate task-start speech and mid-episode interjection/speech pairs."""
-
-    vlm: VlmClient
-    config: Module2Config
-    seed: int = 1729
-    frame_provider: FrameProvider = field(default_factory=null_provider)
-
-    @property
-    def enabled(self) -> bool:
-        return self.config.enabled
-
-    def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
-        rows: list[dict[str, Any]] = []
-        if record.frame_timestamps:
-            t0 = float(record.frame_timestamps[0])
-            initial = self._initial_speech(record)
-            if initial:
-                rows.append(speech_atom(t0, initial))
-        # Pull Module 1's subtask spans for this episode so the
-        # interjection prompt can ground itself in the actual current
-        # subtask at each chosen timestamp. Module 1 ran first.
-        subtask_spans = self._read_subtask_spans(staging)
-        rows.extend(self._mid_episode_interjections(record, subtask_spans))
-        staging.write("module_2", rows)
-
-    @staticmethod
-    def _read_subtask_spans(staging: EpisodeStaging) -> list[dict[str, Any]]:
-        rows = [r for r in staging.read("module_1") if r.get("style") == "subtask"]
-        rows.sort(key=lambda r: float(r["timestamp"]))
-        spans: list[dict[str, Any]] = []
-        last_t: float | None = None
-        for r in rows:
-            t = float(r["timestamp"])
-            if last_t is not None and spans:
-                spans[-1]["end"] = t
-            spans.append({"text": r.get("content") or "", "start": t, "end": t})
-            last_t = t
-        return spans
-
-    @staticmethod
-    def _subtask_at(spans: Sequence[dict[str, Any]], t: float) -> str | None:
-        current: str | None = None
-        for span in spans:
-            if float(span["start"]) <= t:
-                current = span.get("text")
-            else:
-                break
-        return current
-
-    def _initial_speech(self, record: EpisodeRecord) -> str | None:
-        prompt = load_prompt("module_2_initial_speech").format(
-            episode_task=record.episode_task,
-        )
-        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
-        result = self.vlm.generate_json([messages])[0]
-        if isinstance(result, dict) and isinstance(result.get("text"), str):
-            text = result["text"].strip()
-            if text:
-                return text
-        return None
-
-    def _mid_episode_interjections(
-        self,
-        record: EpisodeRecord,
-        subtask_spans: Sequence[dict[str, Any]],
-    ) -> list[dict[str, Any]]:
-        """Generate interjections aligned with the actual demo trajectory.
-
-        Teleop data is frozen — the robot already executed every step in
-        the video. A *counterfactual* interjection like "actually skip
-        the wipe" contradicts what then happens in the video, which is
-        what qwen36moe-10/11 surfaced as low-quality interjections.
-
-        Instead, anchor every interjection at a subtask boundary and
-        write it as a natural user request for the *upcoming* subtask.
-        The robot's visible next behavior IS the interjection's effect,
-        so the training signal stays consistent: interjection text →
-        plan refresh → action stream all line up.
-        """
-        if self.config.max_interjections_per_episode <= 0:
-            return []
-        if len(subtask_spans) < 2:
-            # Need at least one transition (subtask 0 → subtask 1).
-            return []
-        # Deterministic per-episode RNG so reruns are stable across SLURM jobs.
-        rng = random.Random(f"{self.seed}:{record.episode_index}:interjection")
-
-        # Boundaries: the start time of every subtask except the first
-        # (which is just t0 and is covered by the initial-task speech atom).
-        boundaries: list[tuple[float, str, str]] = []
-        for i in range(1, len(subtask_spans)):
-            ts = float(subtask_spans[i]["start"])
-            if ts < self.config.interjection_min_t:
-                continue
-            prev_text = (subtask_spans[i - 1].get("text") or "").strip()
-            next_text = (subtask_spans[i].get("text") or "").strip()
-            if not next_text:
-                continue
-            boundaries.append((ts, prev_text, next_text))
-        if not boundaries:
-            return []
-
-        n = min(self.config.max_interjections_per_episode, len(boundaries))
-        chosen = sorted(rng.sample(boundaries, n), key=lambda b: b[0])
-
-        out: list[dict[str, Any]] = []
-        for t, prev_subtask, next_subtask in chosen:
-            t_snap = _snap_to_frame(t, record.frame_timestamps)
-            # Window straddles the boundary so the VLM sees the end of the
-            # previous subtask and the start of the next one — same
-            # conditioning the policy will see at training time.
-            window_ts = self._window_timestamps(t_snap, record.frame_timestamps)
-            prompt = load_prompt("module_2_interjection").format(
-                episode_task=record.episode_task,
-                prev_subtask=prev_subtask or "(starting from initial state)",
-                next_subtask=next_subtask,
-                timestamp=t_snap,
-                window_seconds=self.config.interjection_window_seconds,
-            )
-            images = self.frame_provider.frames_at(record, window_ts)
-            content = [*to_image_blocks(images), {"type": "text", "text": prompt}]
-            messages = [{"role": "user", "content": content}]
-            result = self.vlm.generate_json([messages])[0]
-            if not isinstance(result, dict):
-                continue
-            interjection_text = result.get("interjection")
-            speech_text = result.get("speech")
-            if not isinstance(interjection_text, str) or not interjection_text.strip():
-                continue
-            if not isinstance(speech_text, str) or not speech_text.strip():
-                continue
-            out.append(
-                {
-                    "role": "user",
-                    "content": interjection_text.strip(),
-                    "style": "interjection",
-                    "timestamp": t_snap,
-                    "tool_calls": None,
-                }
-            )
-            out.append(speech_atom(t_snap, speech_text.strip()))
-        return out
-
-    def _window_timestamps(
-        self, t_anchor: float, frame_timestamps: Sequence[float]
-    ) -> list[float]:
-        """Return a small set of frame timestamps centered on ``t_anchor``.
-
-        The window straddles the subtask boundary the interjection sits
-        on: roughly half the frames cover the end of the previous
-        subtask, half cover the start of the next one. The VLM therefore
-        sees BOTH what just finished AND what's about to start, which is
-        the conditioning we need to write a natural "now please do X"
-        request that matches the visible upcoming behavior.
-        """
-        if not frame_timestamps:
-            return [t_anchor]
-        n = max(1, int(self.config.interjection_window_frames))
-        if n == 1:
-            return [t_anchor]
-        window = float(self.config.interjection_window_seconds)
-        step = window / max(1, n - 1)
-        # Center the window on the anchor so half lands before, half after.
-        start_offset = -window / 2.0
-        targets = [t_anchor + start_offset + step * i for i in range(n)]
-        last_ts = float(frame_timestamps[-1])
-        snapped: list[float] = []
-        seen: set[float] = set()
-        for tgt in targets:
-            clamped = min(last_ts, max(0.0, tgt))
-            t = _snap_to_frame(clamped, frame_timestamps)
-            if t not in seen:
-                seen.add(t)
-                snapped.append(t)
-        return snapped or [t_anchor]

src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py

@@ -1,443 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Module 1: subtask decomposition + plan + memory (PERSISTENT styles)."""
-
-from __future__ import annotations
-
-from collections.abc import Sequence
-from dataclasses import dataclass, field
-from typing import Any
-
-from pathlib import Path
-
-from ..config import Module1Config
-from ..frames import (
-    FrameProvider,
-    VideoFrameProvider,
-    episode_clip_path,
-    null_provider,
-    to_video_block,
-    to_video_url_block,
-)
-from ..prompts import load as load_prompt
-from ..reader import EpisodeRecord
-from ..staging import EpisodeStaging
-from ..vlm_client import VlmClient
-
-
-def _snap_to_frame(t: float, frame_timestamps: Sequence[float]) -> float:
-    """Snap an arbitrary float to the nearest exact source frame timestamp."""
-    if not frame_timestamps:
-        return float(t)
-    nearest = min(frame_timestamps, key=lambda f: abs(f - t))
-    return float(nearest)
-
-
-@dataclass
-class PlanSubtasksMemoryModule:
-    """Generate subtask spans, plan, and memory rows.
-
-    All output is persistent (lives in ``language_persistent``):
-
-    - ``subtask`` rows: one per span, stamped at the span's *start* timestamp
-      (snapped to an exact frame).
-    - ``plan`` rows: emitted at ``t=0``; refreshed at every interjection
-      timestamp via :meth:`run_plan_updates` (called by the executor after
-      Module 2 completes).
-    - ``memory`` rows: emitted at each subtask boundary (= subtask start
-      timestamp from the second subtask onward).
-    """
-
-    vlm: VlmClient
-    config: Module1Config
-    frame_provider: FrameProvider = field(default_factory=null_provider)
-
-    @property
-    def enabled(self) -> bool:
-        return self.config.enabled
-
-    def run_episode(self, record: EpisodeRecord, staging: EpisodeStaging) -> None:
-        rows: list[dict[str, Any]] = []
-        # Resolve the task that drives every other Module-1 prompt. May be
-        # the canonical ``record.episode_task`` (default), or a fresh
-        # description derived from the video when the canonical task is
-        # empty / placeholder / forced-off (see Module1Config.derive_task_*).
-        effective_task = self._resolve_effective_task(record)
-        # ``task_aug`` rows at t=0 (role=user), one per rephrasing — the
-        # PR 1 renderer rotates ``${task}`` deterministically through them
-        # so the policy sees diverse phrasings during training.
-        t0 = float(record.frame_timestamps[0]) if record.frame_timestamps else 0.0
-        if self.config.n_task_rephrasings > 0 and effective_task:
-            rephrasings = self._generate_task_rephrasings(
-                effective_task, n=self.config.n_task_rephrasings
-            )
-            # Always include the effective task itself as the first variant
-            # so the rotation is guaranteed to cover the source-of-truth
-            # phrasing, not just synthetic alternatives.
-            seen: set[str] = set()
-            ordered = [effective_task, *rephrasings]
-            for phrasing in ordered:
-                key = phrasing.strip()
-                if not key or key in seen:
-                    continue
-                seen.add(key)
-                rows.append(
-                    {
-                        "role": "user",
-                        "content": key,
-                        "style": "task_aug",
-                        "timestamp": t0,
-                        "tool_calls": None,
-                    }
-                )
-
-        subtask_spans = self._generate_subtasks(record, task=effective_task)
-        # subtask rows
-        for span in subtask_spans:
-            rows.append(
-                {
-                    "role": "assistant",
-                    "content": span["text"],
-                    "style": "subtask",
-                    "timestamp": _snap_to_frame(span["start"], record.frame_timestamps),
-                    "tool_calls": None,
-                }
-            )
-        # plan row at t=0
-        plan_text = self._generate_plan(record, subtask_spans, task=effective_task)
-        if plan_text is not None:
-            rows.append(
-                {
-                    "role": "assistant",
-                    "content": plan_text,
-                    "style": "plan",
-                    "timestamp": float(t0),
-                    "tool_calls": None,
-                }
-            )
-        # memory rows at every subtask boundary except the very first start
-        prior_memory = ""
-        for i, span in enumerate(subtask_spans[1:], start=1):
-            completed = subtask_spans[i - 1]["text"]
-            remaining = [s["text"] for s in subtask_spans[i:]]
-            mem_text = self._generate_memory(
-                record, prior_memory, completed, remaining, task=effective_task
-            )
-            if mem_text:
-                ts = _snap_to_frame(span["start"], record.frame_timestamps)
-                rows.append(
-                    {
-                        "role": "assistant",
-                        "content": mem_text,
-                        "style": "memory",
-                        "timestamp": ts,
-                        "tool_calls": None,
-                    }
-                )
-                prior_memory = mem_text
-        staging.write("module_1", rows)
-
-    # ------------------------------------------------------------------
-    # Task derivation + rephrasings
-    # ------------------------------------------------------------------
-
-    _PLACEHOLDER_TASKS: frozenset[str] = frozenset(
-        {
-            "debug",
-            "test",
-            "tbd",
-            "todo",
-            "n/a",
-            "na",
-            "untitled",
-            "unnamed",
-            "default",
-            "placeholder",
-        }
-    )
-
-    def _resolve_effective_task(self, record: EpisodeRecord) -> str:
-        """Decide which task string drives Module 1 for this episode.
-
-        Returns the user-supplied ``record.episode_task`` unless
-        ``derive_task_from_video`` says otherwise (see config docstring).
-        Falls back gracefully to the canonical task if video derivation
-        fails.
-        """
-        canonical = (record.episode_task or "").strip()
-        mode = (self.config.derive_task_from_video or "off").strip().lower()
-        if mode == "always":
-            derived = self._derive_task_from_video(record)
-            return derived or canonical
-        if mode == "if_short" and self._task_seems_bad(canonical):
-            derived = self._derive_task_from_video(record)
-            if derived:
-                return derived
-        return canonical
-
-    def _task_seems_bad(self, task: str) -> bool:
-        if not task:
-            return True
-        if len(task.split()) < int(self.config.derive_task_min_words):
-            return True
-        if task.lower() in self._PLACEHOLDER_TASKS:
-            return True
-        return False
-
-    def _derive_task_from_video(self, record: EpisodeRecord) -> str | None:
-        """Ask the VLM "what is this video about" with no task hint at all."""
-        prompt = load_prompt("module_1_video_task")
-        video_block = self._episode_video_block(record)
-        content = [*video_block, {"type": "text", "text": prompt}]
-        messages = [{"role": "user", "content": content}]
-        result = self.vlm.generate_json([messages])[0]
-        if isinstance(result, dict) and isinstance(result.get("task"), str):
-            text = result["task"].strip()
-            if text:
-                return text
-        return None
-
-    def _generate_task_rephrasings(self, base_task: str, *, n: int) -> list[str]:
-        """Generate ``n`` text-only paraphrases of ``base_task``."""
-        if n <= 0 or not base_task:
-            return []
-        prompt = load_prompt("module_1_task_rephrasings").format(
-            base_task=base_task, n=n
-        )
-        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
-        result = self.vlm.generate_json([messages])[0]
-        if not isinstance(result, dict):
-            return []
-        raw = result.get("rephrasings")
-        if not isinstance(raw, list):
-            return []
-        out: list[str] = []
-        for item in raw:
-            if isinstance(item, str):
-                cleaned = item.strip().strip('"').strip("'")
-                if cleaned:
-                    out.append(cleaned)
-        return out[:n]
-
-    def _episode_video_block(self, record: EpisodeRecord) -> list[dict[str, Any]]:
-        """Same video block ``_generate_subtasks`` builds — extracted helper."""
-        if not record.frame_timestamps:
-            return []
-        if self.config.use_video_url and isinstance(self.frame_provider, VideoFrameProvider):
-            cache_dir = Path(self.frame_provider.root) / ".annotate_staging" / ".video_clips"
-            clip = episode_clip_path(record, self.frame_provider, cache_dir)
-            return (
-                to_video_url_block(f"file://{clip}", fps=self.config.use_video_url_fps)
-                if clip is not None
-                else []
-            )
-        episode_duration = record.frame_timestamps[-1] - record.frame_timestamps[0]
-        target_count = max(
-            1, int(round(episode_duration * self.config.frames_per_second))
-        )
-        target_count = min(target_count, self.config.max_video_frames)
-        video_frames = self.frame_provider.video_for_episode(record, target_count)
-        return to_video_block(video_frames)
-
-    def run_plan_updates(
-        self,
-        record: EpisodeRecord,
-        staging: EpisodeStaging,
-        interjection_times: Sequence[float],
-        interjection_texts: Sequence[str] | None = None,
-    ) -> None:
-        """Append additional ``plan`` rows at every interjection timestamp.
-
-        Plans refresh ONLY on user interjections — subtask generation
-        runs ~1 Hz at inference, but plan re-emission is event-driven.
-        Now also forwards the interjection's own text into the prompt so
-        the refreshed plan can actually reflect the user's correction
-        (the previous version told the model "an interjection happened"
-        without telling it what the user said).
-        """
-        existing = staging.read("module_1")
-        spans = self._reconstruct_subtasks_from_rows(existing)
-        already_planned: set[float] = {
-            float(r["timestamp"]) for r in existing if r.get("style") == "plan"
-        }
-        new_rows = list(existing)
-
-        texts: list[str | None] = (
-            [None] * len(interjection_times)
-            if interjection_texts is None
-            else [str(t) if t else None for t in interjection_texts]
-        )
-        for raw_t, inter_text in zip(interjection_times, texts):
-            t = _snap_to_frame(raw_t, record.frame_timestamps)
-            if t in already_planned:
-                continue
-            already_planned.add(t)
-            plan_text = self._generate_plan(
-                record, spans, refresh_t=t, interjection=inter_text
-            )
-            if plan_text is not None:
-                new_rows.append(
-                    {
-                        "role": "assistant",
-                        "content": plan_text,
-                        "style": "plan",
-                        "timestamp": t,
-                        "tool_calls": None,
-                    }
-                )
-        staging.write("module_1", new_rows)
-
-    @staticmethod
-    def _reconstruct_subtasks_from_rows(rows: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
-        out = []
-        last_t: float | None = None
-        for row in sorted(
-            (r for r in rows if r.get("style") == "subtask"),
-            key=lambda r: float(r["timestamp"]),
-        ):
-            t = float(row["timestamp"])
-            if last_t is not None:
-                out[-1]["end"] = t
-            out.append({"text": row.get("content") or "", "start": t, "end": t})
-            last_t = t
-        return out
-
-    def _generate_subtasks(
-        self, record: EpisodeRecord, *, task: str | None = None
-    ) -> list[dict[str, Any]]:
-        if record.row_count == 0 or not record.frame_timestamps:
-            return []
-        episode_duration = record.frame_timestamps[-1] - record.frame_timestamps[0]
-        prompt = load_prompt("module_1_subtasks").format(
-            episode_task=(task if task is not None else record.episode_task),
-            min_subtask_seconds=self.config.min_subtask_seconds,
-            max_steps=self.config.plan_max_steps,
-            episode_duration=f"{episode_duration:.3f}",
-        )
-        if self.config.use_video_url and isinstance(self.frame_provider, VideoFrameProvider):
-            cache_dir = Path(self.frame_provider.root) / ".annotate_staging" / ".video_clips"
-            clip = episode_clip_path(record, self.frame_provider, cache_dir)
-            video_block = (
-                to_video_url_block(f"file://{clip}", fps=self.config.use_video_url_fps)
-                if clip is not None
-                else []
-            )
-        else:
-            target_count = max(
-                1,
-                int(round(episode_duration * self.config.frames_per_second)),
-            )
-            target_count = min(target_count, self.config.max_video_frames)
-            video_frames = self.frame_provider.video_for_episode(record, target_count)
-            video_block = to_video_block(video_frames)
-        content = [*video_block, {"type": "text", "text": prompt}]
-        messages = [{"role": "user", "content": content}]
-        result = self.vlm.generate_json([messages])[0]
-        spans = result.get("subtasks") if isinstance(result, dict) else None
-        if not spans:
-            return []
-        # clamp to [t0, t_last] and sort
-        t0 = record.frame_timestamps[0]
-        t_last = record.frame_timestamps[-1]
-        cleaned: list[dict[str, Any]] = []
-        for span in spans:
-            try:
-                start = float(span["start"])
-                end = float(span["end"])
-                text = str(span["text"]).strip()
-            except (KeyError, ValueError, TypeError):
-                continue
-            start = max(t0, min(start, t_last))
-            end = max(t0, min(end, t_last))
-            if end < start:
-                start, end = end, start
-            if not text:
-                continue
-            cleaned.append({"text": text, "start": start, "end": end})
-        cleaned.sort(key=lambda s: s["start"])
-        return cleaned
-
-    def _generate_plan(
-        self,
-        record: EpisodeRecord,
-        subtask_spans: Sequence[dict[str, Any]],
-        *,
-        refresh_t: float | None = None,
-        interjection: str | None = None,
-        task: str | None = None,
-    ) -> str | None:
-        if not subtask_spans:
-            return None
-        subtasks_text = "\n".join(f"- {s['text']}" for s in subtask_spans)
-        prompt = load_prompt("module_1_plan").format(
-            episode_task=(task if task is not None else record.episode_task),
-            subtasks_text=subtasks_text,
-            plan_max_steps=self.config.plan_max_steps,
-        )
-        if refresh_t is not None:
-            # ``current_subtask`` is the span the refresh time falls into,
-            # so the model knows where in the demonstration the planner is
-            # standing when it re-emits.
-            current_subtask = ""
-            for span in subtask_spans:
-                if float(span["start"]) <= refresh_t and (
-                    "end" not in span or float(span["end"]) > refresh_t
-                ):
-                    current_subtask = span.get("text", "")
-                    break
-            if interjection:
-                prompt += (
-                    f"\n\n(Plan refresh at t={refresh_t:.2f}s after a user "
-                    f"interjection: {interjection!r}. Current subtask just "
-                    f"before the interjection: {current_subtask!r}. Update "
-                    f"the plan so it reflects the interjection — drop or "
-                    f"reorder steps as needed; do not just restate.)\n"
-                )
-            else:
-                # Refresh without an interjection text: still tell the model
-                # where in the episode the plan stands so the re-emission
-                # is grounded. Should be rare — plan refreshes are
-                # interjection-driven by design.
-                prompt += (
-                    f"\n\n(Plan refresh at t={refresh_t:.2f}s. Current "
-                    f"subtask: {current_subtask!r}.)\n"
-                )
-        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
-        result = self.vlm.generate_json([messages])[0]
-        if isinstance(result, dict) and isinstance(result.get("plan"), str):
-            return result["plan"].strip()
-        return None
-
-    def _generate_memory(
-        self,
-        record: EpisodeRecord,
-        prior_memory: str,
-        completed: str,
-        remaining: Sequence[str],
-        *,
-        task: str | None = None,
-    ) -> str:
-        prompt = load_prompt("module_1_memory").format(
-            episode_task=(task if task is not None else record.episode_task),
-            prior_memory=prior_memory or "(none)",
-            completed_subtask=completed,
-            remaining_subtasks=", ".join(remaining) if remaining else "(none)",
-        )
-        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
-        result = self.vlm.generate_json([messages])[0]
-        if isinstance(result, dict) and isinstance(result.get("memory"), str):
-            return result["memory"].strip()
-        return ""

src/lerobot/annotations/steerable_pipeline/prompts/__init__.py

@@ -1,33 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Prompt templates loaded as plain text.
-
-One file per use site. Templates use ``str.format(**vars)`` substitution; we
-intentionally avoid jinja2 here so the templates remain inspectable in
-plain editors and roundtrip cleanly through ``ruff format``.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-_DIR = Path(__file__).parent
-
-
-def load(name: str) -> str:
-    """Read prompt template ``name.txt`` from the ``prompts/`` directory."""
-    path = _DIR / f"{name}.txt"
-    return path.read_text(encoding="utf-8")

src/lerobot/annotations/steerable_pipeline/prompts/module_1_memory.txt

@@ -1,25 +0,0 @@
-You are updating the robot's compressed semantic memory at the boundary of
-a completed subtask.
-
-Reference (verbatim from MEM, Torne 2026):
-"Remove or compress information in the language memory whenever
-appropriate. Keep ONLY the minimal set of relevant information for future
-task execution. Specific object attributes (colors, precise quantities of
-each item) get discarded when their details won't affect subsequent
-actions. Functional outcomes (where items went, how many) are preserved."
-
-Concrete example from MEM:
-  Before: "I put a light green bowl, a dark blue bowl and a bright yellow
-           bowl into the top right cabinet"
-  After:  "I placed three bowls in the top right cabinet"
-
-Episode task: "{episode_task}"
-Previous memory: {prior_memory}
-Just-completed subtask: "{completed_subtask}"
-Remaining subtasks (for relevance judgement only): {remaining_subtasks}
-
-Update the memory. Drop irrelevant detail. Compress completed steps.
-Keep WHAT happened, drop HOW. Shorter is better.
-
-Output strictly valid JSON:
-  {{ "memory": "<one or two short sentences>" }}

src/lerobot/annotations/steerable_pipeline/prompts/module_1_plan.txt

@@ -1,18 +0,0 @@
-You are the high-level planner for a robot demonstrating: "{episode_task}".
-
-Given the subtask decomposition below, write a concise hierarchical PLAN
-the robot should follow. Format the plan as a numbered list, one line per
-high-level step. The plan describes the full task; subtasks are the atomic
-skills used to execute it.
-
-Subtasks for context:
-{subtasks_text}
-
-Authoring rules:
-- 3 to {plan_max_steps} steps.
-- Each step describes one logical chunk of the task, not one motion.
-- Steps must be in execution order.
-- Plain prose, no JSON, no markdown headers.
-
-Output strictly valid JSON:
-  {{ "plan": "1. ...\n2. ...\n3. ..." }}

src/lerobot/annotations/steerable_pipeline/prompts/module_1_subtasks.txt

@@ -1,33 +0,0 @@
-You are labeling a teleoperated robot demonstration.
-
-The user originally asked: "{episode_task}"
-
-You are shown the entire demonstration as a single video. Watch the
-whole clip, then segment it into a list of consecutive atomic subtasks
-the robot performs.
-
-Authoring rules — based on Hi Robot (Shi 2025) atom granularity and
-Pi0.7 (Physical Intelligence 2025) "how, not what" detail:
-
-- Each subtask is one atomic skill the low-level policy can execute,
-  e.g. "pick up one piece of lettuce", "place the bowl into the box",
-  "move the right arm to the left".
-- Capture HOW the subtask is performed, not only WHAT — e.g. prefer
-  "grasp the handle of the sponge with the left hand" to "pick up the
-  sponge".
-- Subtasks are non-overlapping and cover the full episode in order.
-  Choose the cut points yourself based on what you see in the video
-  (gripper open/close events, contact, regrasps, transitions).
-- Each subtask spans at least {min_subtask_seconds} seconds.
-- Do not exceed {max_steps} subtasks total.
-- Every subtask's [start_time, end_time] must lie within
-  [0.0, {episode_duration}] seconds.
-
-Output strictly valid JSON of shape:
-
-  {{
-    "subtasks": [
-      {{"text": "<how-not-what>", "start": <float>, "end": <float>}},
-      ...
-    ]
-  }}

src/lerobot/annotations/steerable_pipeline/prompts/module_1_task_rephrasings.txt

@@ -1,32 +0,0 @@
-You are generating training data for a Hi Robot-style policy. We need
-{n} alternative phrasings of the same robot task so the policy sees
-diverse user prompts during training instead of the same canonical
-string repeated every frame.
-
-Original task:
-"{base_task}"
-
-Generate exactly {n} alternative phrasings of the same task. Vary:
-
-- formality (casual / polite / curt)
-- verbosity (short imperative vs longer polite request)
-- word choice (synonyms, different verbs)
-- sentence structure (imperative / question / suggestion)
-
-Hard rules:
-- Each phrasing MUST preserve the exact meaning of the original task.
-  Do not change which object is involved, the destination, or the
-  action. Do not add extra steps. Do not invent new objects.
-- Each phrasing must be a single short sentence, plain prose, no
-  markdown, no quotes, no list numbers.
-- Phrasings must be distinct — no near-duplicates.
-- Output exactly {n} entries.
-
-Output strictly valid JSON:
-  {{
-    "rephrasings": [
-      "<phrasing 1>",
-      "<phrasing 2>",
-      ...
-    ]
-  }}

src/lerobot/annotations/steerable_pipeline/prompts/module_1_video_task.txt

@@ -1,17 +0,0 @@
-The video above shows a robot manipulation episode in full. Look at
-the entire video and describe in ONE concise sentence what the robot
-is doing.
-
-Rules:
-- One sentence, in natural English, like a user instruction.
-- Capture the goal of the demonstration, not low-level motions.
-  Example: "place the yellow cube into the red bin" — not "move the
-  end-effector down 5cm and close the gripper".
-- 4 to 15 words. Plain prose, no markdown, no bullets, no quotes.
-- Do not invent objects or actions that aren't visible.
-- Do not output anything other than the JSON object below.
-
-Output strictly valid JSON:
-  {{
-    "task": "<single concise sentence describing what the robot does in this video>"
-  }}

src/lerobot/annotations/steerable_pipeline/prompts/module_2_initial_speech.txt

@@ -1,10 +0,0 @@
-The user just asked the robot: "{episode_task}".
-
-Generate a short verbal acknowledgement the robot would speak back before
-beginning the task. Style: confident, friendly, single short sentence.
-
-Examples (Hi Robot, Shi 2025): "Sure, I won't put cheese on it.",
-"OK, starting with the sponge.", "Got it.".
-
-Output strictly valid JSON:
-  {{ "text": "<the spoken acknowledgement>" }}

src/lerobot/annotations/steerable_pipeline/prompts/module_2_interjection.txt

@@ -1,46 +0,0 @@
-You are generating training data for a Hi Robot-style hierarchical
-robot policy. The robot in this demonstration has ALREADY executed
-every step shown in the video — we cannot retroactively change the
-action stream. To keep training data consistent with the video, the
-"interjection" must align with what the robot is *about to do next* in
-the demonstration, framed as a natural mid-task user request.
-
-The episode's overall task: "{episode_task}".
-
-The images above show roughly {window_seconds:.1f} seconds straddling a
-subtask boundary in the demonstration:
-
-- Subtask the robot just finished: "{prev_subtask}"
-- Subtask the robot is about to start: "{next_subtask}"
-- Time into episode: {timestamp:.2f}s
-
-Write ONE interjection the user would naturally say at this moment to
-prompt / confirm / encourage the robot to do "{next_subtask}". Phrase it
-like a real human mid-task remark — conversational, varied, sometimes
-just a nudge, sometimes a clarification, sometimes a small constraint
-that the upcoming motion happens to satisfy. Plus the robot's verbal
-acknowledgement.
-
-Hard rules:
-
-- The interjection MUST be consistent with the next subtask. The user
-  cannot ask for something different from what the robot then does in
-  the video. If you're tempted to say "actually skip X" or "do Y
-  instead", DO NOT — those would contradict the demonstration.
-- The interjection must reference an object, location, or action that
-  is plausible given the visible scene and the next subtask text.
-- One sentence each. Conversational, not robotic.
-
-Style examples (vary the phrasing — don't reuse these verbatim):
-  - "Now go ahead and {next_subtask}."
-  - "Great, can you {next_subtask} next?"
-  - "{next_subtask}, please."
-  - "Before you continue, please {next_subtask}."
-  - "Looking good — {next_subtask} now."
-  - "Okay, {next_subtask}."
-
-Output strictly valid JSON:
-  {{
-    "interjection": "<single sentence the user says, asking for the next subtask>",
-    "speech":       "<single sentence the robot speaks back, confirming and starting>"
-  }}

src/lerobot/annotations/steerable_pipeline/prompts/module_3_vqa.txt

@@ -1,32 +0,0 @@
-You are generating a frame-grounded visual question/answer pair for
-chain-of-thought training. Reference: ECoT (Zawalski 2024) and Steerable
-Policies — both train policies on grounded features such as bounding box
-pixel coordinates, keypoints, counts, attributes, and spatial relations.
-
-The frame shows a robot working on: "{episode_task}".
-
-Question types and the EXACT answer JSON shape required for each:
-
-  bbox       => {{"detections": [{{"label": "<obj>", "bbox_format": "xyxy",
-                                    "bbox": [x1, y1, x2, y2]}}, ...]}}
-                bbox is in pixel coordinates (x_min, y_min, x_max, y_max).
-                ECoT example: "a white cup [124, 25, 176, 113]".
-
-  keypoint   => {{"label": "<point>", "point_format": "xy",
-                  "point": [x, y]}}
-
-  count      => {{"label": "<obj>", "count": <int>,
-                  "note": "<optional short note>"}}
-
-  attribute  => {{"label": "<obj>", "attribute": "<color|shape|state|...>",
-                  "value": "<observed value>"}}
-
-  spatial    => {{"subject": "<obj>", "relation": "<left_of|right_of|on|in|"
-                  "above|below|near>", "object": "<obj>"}}
-
-Generate a question of type "{question_type}". Output strictly valid JSON:
-
-  {{
-    "question": "<short, frame-grounded question>",
-    "answer":   <object whose shape matches the schema above>
-  }}

src/lerobot/annotations/steerable_pipeline/reader.py

@@ -1,219 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Datatrove-shaped reader.
-
-The reader walks ``data/chunk-*/file-*.parquet`` and yields one record per
-episode containing:
-
-- ``episode_index``: int
-- ``frame_timestamps``: tuple[float, ...]
-- ``frame_indices``: tuple[int, ...]
-- ``episode_task``: str (canonical task from ``meta/tasks.parquet``)
-- ``data_path``: pathlib.Path of the source parquet shard
-- ``frames_df``: pandas.DataFrame slice for the episode (only loaded on demand)
-
-This shape lets each module operate per-episode without loading all parquet
-rows into memory at once. It deliberately does not depend on datatrove —
-datatrove integration wraps this generator inside a ``PipelineStep`` in
-:mod:`.executor`.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Iterator
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import pyarrow.parquet as pq
-
-from lerobot.datasets.utils import DEFAULT_TASKS_PATH
-
-
-@dataclass
-class EpisodeRecord:
-    """Per-episode record yielded by the reader."""
-
-    episode_index: int
-    episode_task: str
-    frame_timestamps: tuple[float, ...]
-    frame_indices: tuple[int, ...]
-    data_path: Path
-    row_offset: int  # row offset within the parquet file where this episode starts
-    row_count: int  # number of rows for this episode
-
-    def frames_df(self):  # type: ignore[no-untyped-def]
-        """Lazy-load the pandas slice for this episode."""
-        import pandas as pd  # noqa: PLC0415  - deferred for optional dataset extra
-
-        table = pq.read_table(self.data_path)
-        df: pd.DataFrame = table.to_pandas()
-        slice_ = df.iloc[self.row_offset : self.row_offset + self.row_count].reset_index(drop=True)
-        return slice_
-
-
-def _load_tasks_lookup(root: Path) -> dict[int, str]:
-    tasks_path = root / DEFAULT_TASKS_PATH
-    if not tasks_path.exists():
-        return {}
-    table = pq.read_table(tasks_path)
-    cols = {name: table.column(name).to_pylist() for name in table.column_names}
-    if "task_index" in cols and "task" in cols:
-        return dict(zip(cols["task_index"], cols["task"], strict=True))
-    raise ValueError(f"meta/tasks.parquet at {tasks_path} missing 'task_index' or 'task'")
-
-
-def iter_episodes(root: Path, *, only_episodes: tuple[int, ...] | None = None) -> Iterator[EpisodeRecord]:
-    """Yield :class:`EpisodeRecord` for every episode under ``root/data/``.
-
-    Episodes are yielded in ascending ``episode_index`` order. The reader does
-    not assume a specific chunk/file layout: it scans every ``*.parquet``
-    under ``data/`` and groups by ``episode_index``.
-    """
-    tasks = _load_tasks_lookup(root)
-    data_dir = root / "data"
-    parquet_files = sorted(data_dir.rglob("*.parquet"))
-
-    only_set = set(only_episodes) if only_episodes is not None else None
-
-    for path in parquet_files:
-        yield from _iter_one_path(path, tasks, only_set)
-
-
-def _iter_one_path(path: Path, tasks: dict[int, str], only_set: set[int] | None) -> Iterator[EpisodeRecord]:
-    table = pq.read_table(path)
-    names = table.column_names
-    if "episode_index" not in names:
-        return
-    episode_col = table.column("episode_index").to_pylist()
-    timestamp_col = (
-        table.column("timestamp").to_pylist() if "timestamp" in names else [0.0] * len(episode_col)
-    )
-    frame_col = (
-        table.column("frame_index").to_pylist() if "frame_index" in names else list(range(len(episode_col)))
-    )
-    task_col = table.column("task_index").to_pylist() if "task_index" in names else None
-
-    def _build(
-        ep: int,
-        start: int,
-        end: int,
-        task_idx: int | None,
-        ts_buf: list[float],
-        fi_buf: list[int],
-    ) -> EpisodeRecord | None:
-        if only_set is not None and ep not in only_set:
-            return None
-        task = tasks.get(task_idx, "") if task_idx is not None else ""
-        return EpisodeRecord(
-            episode_index=ep,
-            episode_task=task,
-            frame_timestamps=tuple(ts_buf),
-            frame_indices=tuple(fi_buf),
-            data_path=path,
-            row_offset=start,
-            row_count=end - start,
-        )
-
-    cur_ep: int | None = None
-    start_offset = 0
-    ts_buf: list[float] = []
-    fi_buf: list[int] = []
-    cur_task_idx: int | None = None
-
-    for i, ep in enumerate(episode_col):
-        if cur_ep is None:
-            cur_ep = ep
-            start_offset = i
-            ts_buf = [timestamp_col[i]]
-            fi_buf = [frame_col[i]]
-            cur_task_idx = task_col[i] if task_col is not None else None
-            continue
-        if ep != cur_ep:
-            rec = _build(cur_ep, start_offset, i, cur_task_idx, ts_buf, fi_buf)
-            if rec is not None:
-                yield rec
-            cur_ep = ep
-            start_offset = i
-            ts_buf = [timestamp_col[i]]
-            fi_buf = [frame_col[i]]
-            cur_task_idx = task_col[i] if task_col is not None else None
-        else:
-            ts_buf.append(timestamp_col[i])
-            fi_buf.append(frame_col[i])
-
-    if cur_ep is not None:
-        rec = _build(cur_ep, start_offset, len(episode_col), cur_task_idx, ts_buf, fi_buf)
-        if rec is not None:
-            yield rec
-
-
-def gather_data_paths(root: Path) -> list[Path]:
-    """Return every ``data/chunk-*/file-*.parquet`` path under ``root``."""
-    return sorted((root / "data").rglob("*.parquet"))
-
-
-def episode_offsets_per_path(path: Path) -> dict[int, tuple[int, int]]:
-    """Return ``{episode_index: (row_offset, row_count)}`` for one parquet."""
-    table = pq.read_table(path, columns=["episode_index"])
-    episode_col = table.column("episode_index").to_pylist()
-    out: dict[int, tuple[int, int]] = {}
-    cur_ep: int | None = None
-    start = 0
-    for i, ep in enumerate(episode_col):
-        if cur_ep is None:
-            cur_ep = ep
-            start = i
-            continue
-        if ep != cur_ep:
-            out[cur_ep] = (start, i - start)
-            cur_ep = ep
-            start = i
-    if cur_ep is not None:
-        out[cur_ep] = (start, len(episode_col) - start)
-    return out
-
-
-def keyframe_indices(record: EpisodeRecord, k: int) -> list[int]:
-    """Return ``k`` evenly spaced row indices into the episode (relative)."""
-    n = record.row_count
-    if k <= 0 or n == 0:
-        return []
-    if k >= n:
-        return list(range(n))
-    step = (n - 1) / (k - 1) if k > 1 else 0.0
-    return [int(round(i * step)) for i in range(k)] if k > 1 else [n // 2]
-
-
-def lookup_data_path(root: Path, episode_index: int) -> tuple[Path, int, int] | None:
-    """Find the parquet file containing ``episode_index`` and its slice bounds."""
-    for path in gather_data_paths(root):
-        offsets = episode_offsets_per_path(path)
-        if episode_index in offsets:
-            start, count = offsets[episode_index]
-            return path, start, count
-    return None
-
-
-def episode_frame_timestamps(root: Path, episode_index: int) -> tuple[Any, list[float]]:
-    """Return the parquet path and per-frame timestamps for ``episode_index``."""
-    found = lookup_data_path(root, episode_index)
-    if found is None:
-        raise ValueError(f"Episode {episode_index} not found under {root}/data/")
-    path, start, count = found
-    table = pq.read_table(path, columns=["timestamp"])
-    timestamps = table.column("timestamp").to_pylist()[start : start + count]
-    return path, [float(t) for t in timestamps]

src/lerobot/annotations/steerable_pipeline/staging.py

@@ -1,98 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Per-episode staging.
-
-Each module writes its raw output as a JSONL file under
-``<staging_dir>/episode_{ep:06d}/<module>.jsonl``. The writer reads back this
-staging tree and partitions rows into the two language columns.
-
-JSONL is preferred over parquet here because the staging artifact is meant to
-be human-inspectable, easy to diff between prompt iterations, and trivially
-appended to. The final dataset format is parquet; staging is just an
-intermediate.
-"""
-
-from __future__ import annotations
-
-import json
-from collections.abc import Iterable, Iterator
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-ModuleName = str
-
-_MODULES: tuple[ModuleName, ...] = (
-    "module_1",
-    "module_2",
-    "module_3",
-)
-
-
-@dataclass
-class EpisodeStaging:
-    """Filesystem layout for a single episode's staged module outputs."""
-
-    root: Path
-    episode_index: int
-
-    @property
-    def episode_dir(self) -> Path:
-        return self.root / f"episode_{self.episode_index:06d}"
-
-    def path_for(self, module: ModuleName) -> Path:
-        if module not in _MODULES:
-            raise ValueError(f"Unknown module {module!r}; expected one of {_MODULES}")
-        return self.episode_dir / f"{module}.jsonl"
-
-    def write(self, module: ModuleName, rows: Iterable[dict[str, Any]]) -> Path:
-        path = self.path_for(module)
-        path.parent.mkdir(parents=True, exist_ok=True)
-        with path.open("w", encoding="utf-8") as f:
-            for row in rows:
-                f.write(json.dumps(row, ensure_ascii=False, sort_keys=True))
-                f.write("\n")
-        return path
-
-    def read(self, module: ModuleName) -> list[dict[str, Any]]:
-        path = self.path_for(module)
-        if not path.exists():
-            return []
-        out: list[dict[str, Any]] = []
-        with path.open(encoding="utf-8") as f:
-            for line in f:
-                line = line.strip()
-                if line:
-                    out.append(json.loads(line))
-        return out
-
-    def read_all(self) -> dict[ModuleName, list[dict[str, Any]]]:
-        return {m: self.read(m) for m in _MODULES}
-
-    def has(self, module: ModuleName) -> bool:
-        return self.path_for(module).exists()
-
-
-def iter_staged_episodes(root: Path) -> Iterator[int]:
-    """Yield episode indices for which any staging artifact exists."""
-    if not root.exists():
-        return
-    for child in sorted(root.iterdir()):
-        if child.is_dir() and child.name.startswith("episode_"):
-            try:
-                yield int(child.name.removeprefix("episode_"))
-            except ValueError:
-                continue

src/lerobot/annotations/steerable_pipeline/validator.py

@@ -1,334 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Pre-write validation against staged outputs.
-
-Runs after Modules 1–3 have all written their per-episode artifacts but
-*before* the writer rewrites parquet shards. The validator never touches
-parquet; it only inspects the staging tree and the source frame timestamps
-exposed by :class:`EpisodeRecord`.
-
-Checks (per the plan's "Intermediate staging and validation" section):
-
-- exact timestamp alignment against source frame timestamps
-- no orphan speech / interjection pairs
-- plan / memory emission consistency (events have a paired persistent row)
-- VQA assistant ``content`` is valid JSON (one of bbox / keypoint / count /
-  attribute / spatial)
-- every row maps to its correct column under :func:`column_for_style`
-"""
-
-from __future__ import annotations
-
-import json
-import logging
-from collections.abc import Iterable, Sequence
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any
-
-from lerobot.datasets.language import (
-    LANGUAGE_EVENTS,
-    LANGUAGE_PERSISTENT,
-    column_for_style,
-    is_view_dependent_style,
-    validate_camera_field,
-)
-
-from .reader import EpisodeRecord
-from .staging import EpisodeStaging
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class ValidationReport:
-    """Outcome of one validation pass across all episodes."""
-
-    errors: list[str] = field(default_factory=list)
-    warnings: list[str] = field(default_factory=list)
-    episodes_checked: int = 0
-
-    @property
-    def ok(self) -> bool:
-        return not self.errors
-
-    def add_error(self, message: str) -> None:
-        self.errors.append(message)
-
-    def add_warning(self, message: str) -> None:
-        self.warnings.append(message)
-
-    def summary(self) -> str:
-        return f"checked={self.episodes_checked} errors={len(self.errors)} warnings={len(self.warnings)}"
-
-
-VQA_ANSWER_SHAPES: dict[str, set[str]] = {
-    "bbox": {"detections"},
-    "keypoint": {"label", "point_format", "point"},
-    "count": {"label", "count"},
-    "attribute": {"label", "attribute", "value"},
-    "spatial": {"subject", "relation", "object"},
-}
-
-
-def classify_vqa_answer(payload: Any) -> str | None:
-    """Best-effort classification of a VQA answer payload to a question type."""
-    if not isinstance(payload, dict):
-        return None
-    keys = set(payload.keys())
-    for kind, required in VQA_ANSWER_SHAPES.items():
-        if required.issubset(keys):
-            return kind
-    return None
-
-
-@dataclass
-class StagingValidator:
-    """Walks the staging tree and produces a :class:`ValidationReport`."""
-
-    timestamp_atol: float = 0.0  # exact-match by default
-    dataset_camera_keys: tuple[str, ...] | None = None
-    """Known ``observation.images.*`` keys on the dataset. When set, the
-    validator additionally enforces that every view-dependent row's
-    ``camera`` field references one of these keys. Pass ``None`` (default)
-    to skip that cross-check (e.g. in unit tests with no real dataset)."""
-
-    def validate(
-        self,
-        records: Sequence[EpisodeRecord],
-        staging_dir: Path,
-    ) -> ValidationReport:
-        report = ValidationReport()
-        for record in records:
-            self._validate_episode(record, staging_dir, report)
-            report.episodes_checked += 1
-        return report
-
-    def _validate_episode(
-        self,
-        record: EpisodeRecord,
-        staging_dir: Path,
-        report: ValidationReport,
-    ) -> None:
-        staging = EpisodeStaging(staging_dir, record.episode_index)
-        staged = staging.read_all()
-        all_rows: list[dict[str, Any]] = []
-        for module_name, rows in staged.items():
-            for row in rows:
-                row = {**row, "_module": module_name}
-                all_rows.append(row)
-
-        frame_ts = set(record.frame_timestamps)
-
-        events: list[dict[str, Any]] = []
-        persistent: list[dict[str, Any]] = []
-        for row in all_rows:
-            self._check_column_routing(row, report, record.episode_index)
-            self._check_camera_field(
-                row, report, record.episode_index, self.dataset_camera_keys
-            )
-            if column_for_style(row.get("style")) == LANGUAGE_PERSISTENT:
-                persistent.append(row)
-            else:
-                events.append(row)
-
-        for row in events:
-            self._check_event_timestamp_alignment(row, frame_ts, report, record.episode_index)
-
-        self._check_speech_interjection_pairs(events, report, record.episode_index)
-        self._check_plan_memory_consistency(persistent, events, report, record.episode_index)
-        self._check_vqa_json(events, report, record.episode_index)
-        self._check_vqa_uniqueness_per_frame_camera(events, report, record.episode_index)
-
-    def _check_camera_field(
-        self,
-        row: dict[str, Any],
-        report: ValidationReport,
-        episode_index: int,
-        dataset_camera_keys: Sequence[str] | None,
-    ) -> None:
-        """Enforce the camera invariant + that the key matches the dataset's cameras."""
-        style = row.get("style")
-        camera = row.get("camera")
-        try:
-            validate_camera_field(style, camera)
-        except ValueError as exc:
-            report.add_error(
-                f"ep={episode_index} module={row.get('_module')}: {exc}"
-            )
-            return
-        if (
-            is_view_dependent_style(style)
-            and dataset_camera_keys
-            and camera not in dataset_camera_keys
-        ):
-            report.add_error(
-                f"ep={episode_index} module={row.get('_module')}: camera {camera!r} on style "
-                f"{style!r} is not one of the dataset's video keys {sorted(dataset_camera_keys)!r}"
-            )
-
-    def _check_vqa_uniqueness_per_frame_camera(
-        self,
-        events: Iterable[dict[str, Any]],
-        report: ValidationReport,
-        episode_index: int,
-    ) -> None:
-        """Ensure at most one (vqa, user) and one (vqa, assistant) per (t, camera)."""
-        counts: dict[tuple[float, str, str], int] = {}
-        for row in events:
-            if row.get("style") != "vqa":
-                continue
-            ts = row.get("timestamp")
-            camera = row.get("camera")
-            role = row.get("role")
-            if ts is None or camera is None or role is None:
-                continue  # other validators flag these
-            key = (float(ts), str(camera), str(role))
-            counts[key] = counts.get(key, 0) + 1
-        for (ts, camera, role), n in counts.items():
-            if n > 1:
-                report.add_error(
-                    f"ep={episode_index}: {n} duplicate vqa rows at t={ts} "
-                    f"camera={camera!r} role={role!r}; expected at most one per (t, camera, role)"
-                )
-
-    def _check_column_routing(
-        self,
-        row: dict[str, Any],
-        report: ValidationReport,
-        episode_index: int,
-    ) -> None:
-        style = row.get("style")
-        module = row.get("_module")
-        try:
-            target_col = column_for_style(style)
-        except ValueError:
-            report.add_error(f"ep={episode_index} module={module}: unknown style {style!r}")
-            return
-        if module == "module_1" and target_col != LANGUAGE_PERSISTENT:
-            report.add_error(
-                f"ep={episode_index} module=module_1 emitted style {style!r} that routes to {target_col} (must be persistent)"
-            )
-        if module in {"module_2", "module_3"} and target_col != LANGUAGE_EVENTS:
-            report.add_error(
-                f"ep={episode_index} module={module} emitted style {style!r} that routes to {target_col} (must be events)"
-            )
-
-    def _check_event_timestamp_alignment(
-        self,
-        row: dict[str, Any],
-        frame_ts: set[float],
-        report: ValidationReport,
-        episode_index: int,
-    ) -> None:
-        ts = row.get("timestamp")
-        if ts is None:
-            report.add_error(f"ep={episode_index}: event row missing timestamp: {row!r}")
-            return
-        if self.timestamp_atol == 0.0:
-            if float(ts) not in frame_ts:
-                report.add_error(
-                    f"ep={episode_index}: event row timestamp {ts!r} does not match any source frame timestamp"
-                )
-        else:
-            if not any(abs(float(ts) - f) <= self.timestamp_atol for f in frame_ts):
-                report.add_error(
-                    f"ep={episode_index}: event row timestamp {ts!r} not within {self.timestamp_atol}s of any frame"
-                )
-
-    def _check_speech_interjection_pairs(
-        self,
-        events: Iterable[dict[str, Any]],
-        report: ValidationReport,
-        episode_index: int,
-    ) -> None:
-        speech_ts: dict[float, int] = {}
-        interjection_ts: dict[float, int] = {}
-        for row in events:
-            ts = row.get("timestamp")
-            if ts is None:
-                continue
-            ts_f = float(ts)
-            if row.get("style") is None and row.get("role") == "assistant":
-                speech_ts[ts_f] = speech_ts.get(ts_f, 0) + 1
-            if row.get("style") == "interjection":
-                interjection_ts[ts_f] = interjection_ts.get(ts_f, 0) + 1
-
-        for ts in interjection_ts:
-            if ts not in speech_ts:
-                report.add_error(f"ep={episode_index}: interjection at t={ts} has no paired speech atom")
-
-    def _check_plan_memory_consistency(
-        self,
-        persistent: Sequence[dict[str, Any]],
-        events: Sequence[dict[str, Any]],
-        report: ValidationReport,
-        episode_index: int,
-    ) -> None:
-        plan_ts = sorted({float(r["timestamp"]) for r in persistent if r.get("style") == "plan"})
-        memory_ts = sorted({float(r["timestamp"]) for r in persistent if r.get("style") == "memory"})
-        subtask_ts = sorted({float(r["timestamp"]) for r in persistent if r.get("style") == "subtask"})
-        interjection_ts = sorted(
-            {
-                float(r["timestamp"])
-                for r in events
-                if r.get("style") == "interjection" and r.get("timestamp") is not None
-            }
-        )
-
-        if persistent and not plan_ts:
-            report.add_warning(f"ep={episode_index}: persistent rows present but no plan emitted")
-        # every interjection should have a same-timestamp plan refresh
-        for ts in interjection_ts:
-            if ts not in set(plan_ts):
-                report.add_error(
-                    f"ep={episode_index}: interjection at t={ts} has no co-timestamped plan update"
-                )
-        # memory should be emitted at subtask boundaries (subset relation)
-        if memory_ts and subtask_ts:
-            mem_set = set(memory_ts)
-            sub_set = set(subtask_ts)
-            stray = sorted(mem_set - sub_set)
-            if stray:
-                report.add_warning(f"ep={episode_index}: memory rows at {stray} not at any subtask boundary")
-
-    def _check_vqa_json(
-        self,
-        events: Iterable[dict[str, Any]],
-        report: ValidationReport,
-        episode_index: int,
-    ) -> None:
-        for row in events:
-            if row.get("style") != "vqa" or row.get("role") != "assistant":
-                continue
-            content = row.get("content")
-            if content is None:
-                report.add_error(
-                    f"ep={episode_index}: VQA assistant row at t={row.get('timestamp')} has null content"
-                )
-                continue
-            try:
-                payload = json.loads(content)
-            except (TypeError, ValueError) as exc:
-                report.add_error(
-                    f"ep={episode_index}: VQA assistant content not valid JSON at t={row.get('timestamp')}: {exc}"
-                )
-                continue
-            shape = classify_vqa_answer(payload)
-            if shape is None:
-                report.add_error(
-                    f"ep={episode_index}: VQA assistant payload at t={row.get('timestamp')} does not match any known shape: keys={list(payload) if isinstance(payload, dict) else type(payload).__name__}"
-                )

src/lerobot/annotations/steerable_pipeline/vlm_client.py

@@ -1,741 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Shared Qwen-VL client.
-
-The pipeline uses a single shared VLM across modules. vLLM is preferred when
-available (high throughput, JSON-guided decoding); transformers is the
-fallback. A ``stub`` backend is used for unit tests so fixtures never call
-into a real model.
-
-The client speaks one method, :meth:`VlmClient.generate_json`, which:
-
-- accepts a list of OpenAI/HF-style multimodal messages,
-- requests JSON output (``json_mode=True`` enables guided decoding when the
-  backend supports it),
-- batches requests transparently,
-- and reprompts once on a JSON parse failure with an inline correction
-  message before raising.
-"""
-
-from __future__ import annotations
-
-import json
-import os
-import threading
-from collections.abc import Callable, Sequence
-from dataclasses import dataclass
-from typing import Any, Protocol
-
-from .config import VlmConfig
-
-
-class VlmClient(Protocol):
-    """Protocol every backend must implement."""
-
-    def generate_json(
-        self,
-        messages_batch: Sequence[Sequence[dict[str, Any]]],
-        *,
-        max_new_tokens: int | None = None,
-        temperature: float | None = None,
-    ) -> list[Any]:
-        """Generate one JSON-decoded response per messages list."""
-
-
-@dataclass
-class StubVlmClient:
-    """Deterministic stub used in unit tests.
-
-    A test passes a callable that maps the *last user message text* (or, if
-    that is empty, the full message list) to a JSON-serializable response.
-    """
-
-    responder: Callable[[Sequence[dict[str, Any]]], Any]
-
-    def generate_json(
-        self,
-        messages_batch: Sequence[Sequence[dict[str, Any]]],
-        *,
-        max_new_tokens: int | None = None,
-        temperature: float | None = None,
-    ) -> list[Any]:
-        return [self.responder(list(messages)) for messages in messages_batch]
-
-
-def _strip_to_json(text: str) -> Any:
-    text = text.strip()
-    # Strip <think>...</think> blocks (Qwen3 Thinking style)
-    while "<think>" in text and "</think>" in text:
-        start = text.find("<think>")
-        end = text.find("</think>", start) + len("</think>")
-        text = (text[:start] + text[end:]).strip()
-    # Strip ```json ... ``` fences from chat-tuned backbones
-    if text.startswith("```"):
-        first = text.find("\n")
-        last = text.rfind("```")
-        if first != -1 and last != -1 and last > first:
-            text = text[first + 1 : last].strip()
-    try:
-        return json.loads(text)
-    except (ValueError, json.JSONDecodeError):
-        pass
-    # Fall back to extracting the first balanced {...} block.
-    obj_text = _extract_first_json_object(text)
-    if obj_text is None:
-        raise json.JSONDecodeError("No JSON object found", text, 0)
-    return json.loads(obj_text)
-
-
-def _extract_first_json_object(text: str) -> str | None:
-    """Return the first balanced ``{...}`` substring, ignoring braces in
-    string literals. Returns ``None`` if no balanced block is found."""
-    start = text.find("{")
-    if start < 0:
-        return None
-    depth = 0
-    in_string = False
-    escape = False
-    for i in range(start, len(text)):
-        ch = text[i]
-        if escape:
-            escape = False
-            continue
-        if ch == "\\":
-            escape = True
-            continue
-        if ch == '"' and not escape:
-            in_string = not in_string
-            continue
-        if in_string:
-            continue
-        if ch == "{":
-            depth += 1
-        elif ch == "}":
-            depth -= 1
-            if depth == 0:
-                return text[start : i + 1]
-    return None
-
-
-@dataclass
-class _GenericTextClient:
-    """Wraps any text-generation callable in JSON-mode + one-retry semantics."""
-
-    generate_text: Callable[[Sequence[Sequence[dict[str, Any]]], int, float], list[str]]
-    config: VlmConfig
-
-    def generate_json(
-        self,
-        messages_batch: Sequence[Sequence[dict[str, Any]]],
-        *,
-        max_new_tokens: int | None = None,
-        temperature: float | None = None,
-    ) -> list[Any]:
-        max_tok = max_new_tokens if max_new_tokens is not None else self.config.max_new_tokens
-        temp = temperature if temperature is not None else self.config.temperature
-        raw = self.generate_text(messages_batch, max_tok, temp)
-        out: list[Any] = []
-        for messages, text in zip(messages_batch, raw, strict=True):
-            try:
-                out.append(_strip_to_json(text))
-                continue
-            except (ValueError, json.JSONDecodeError):
-                pass
-            retry = list(messages) + [
-                {"role": "assistant", "content": text},
-                {
-                    "role": "user",
-                    "content": (
-                        "Your previous reply was not valid JSON. "
-                        "Reply with strictly valid JSON, no prose, no fences."
-                    ),
-                },
-            ]
-            retry_text = self.generate_text([retry], max_tok, temp)[0]
-            try:
-                out.append(_strip_to_json(retry_text))
-            except (ValueError, json.JSONDecodeError):
-                # After retry: log preview and return None instead of crashing
-                # the whole pipeline. Modules treat None as "skip".
-                preview = retry_text.strip().replace("\n", " ")[:200]
-                print(
-                    f"[vlm] WARNING: failed to parse JSON after retry; preview: {preview!r}",
-                    flush=True,
-                )
-                out.append(None)
-        return out
-
-
-def make_vlm_client(config: VlmConfig) -> VlmClient:
-    """Build the shared VLM client per the configured backend.
-
-    For ``stub``, callers should construct :class:`StubVlmClient` directly with
-    a responder callable. ``stub`` here is rejected to make accidental misuse
-    obvious.
-    """
-    if config.backend == "stub":
-        raise ValueError(
-            "Use StubVlmClient(...) directly for the stub backend; make_vlm_client builds real clients."
-        )
-    if config.backend == "vllm":
-        return _make_vllm_client(config)
-    if config.backend == "transformers":
-        return _make_transformers_client(config)
-    if config.backend == "openai":
-        return _make_openai_client(config)
-    raise ValueError(f"Unknown VLM backend: {config.backend!r}")
-
-
-def _make_vllm_client(config: VlmConfig) -> VlmClient:
-    try:
-        from vllm import LLM, SamplingParams  # type: ignore[import-not-found]
-    except ImportError as exc:
-        raise ImportError(
-            "vllm is required for backend='vllm'. Install with `pip install lerobot[annotations]`."
-        ) from exc
-    # Workaround for cuDNN 9.x + torch 2.8 conv3d regression that surfaces
-    # as CUDNN_STATUS_NOT_INITIALIZED in Qwen-VL vision-tower patch
-    # embedders. Setting LEROBOT_DISABLE_CUDNN=1 forces native PyTorch
-    # convolution kernels — slower but functional.
-    import os as _os  # noqa: PLC0415
-
-    if _os.environ.get("LEROBOT_DISABLE_CUDNN", "").lower() in {"1", "true", "yes"}:
-        import torch as _torch  # noqa: PLC0415
-
-        _torch.backends.cudnn.enabled = False
-    llm_kwargs: dict[str, Any] = {
-        "model": config.model_id,
-        "tensor_parallel_size": config.tensor_parallel_size,
-        "gpu_memory_utilization": config.gpu_memory_utilization,
-        "trust_remote_code": config.trust_remote_code,
-    }
-    if config.max_model_len is not None:
-        llm_kwargs["max_model_len"] = config.max_model_len
-    llm = LLM(**llm_kwargs)
-
-    def _gen(batch: Sequence[Sequence[dict[str, Any]]], max_tok: int, temp: float) -> list[str]:
-        # ``guided_decoding`` would speed up parsing but its API differs across
-        # vllm releases (dict vs GuidedDecodingParams). The _GenericTextClient
-        # wrapper already has a one-retry JSON-recovery path, so we skip it.
-        params = SamplingParams(max_tokens=max_tok, temperature=temp)
-        # ``llm.chat`` handles chat-template application + multimodal input
-        # extraction (image/video blocks) internally, which ``llm.generate``
-        # does not.
-        outputs = llm.chat([list(m) for m in batch], params)
-        return [o.outputs[0].text for o in outputs]
-
-    return _GenericTextClient(_gen, config)
-
-
-def _make_transformers_client(config: VlmConfig) -> VlmClient:
-    try:
-        import torch  # type: ignore[import-not-found]
-        import transformers  # type: ignore[import-not-found]
-        from transformers import AutoProcessor  # type: ignore[import-not-found]
-    except ImportError as exc:
-        raise ImportError("transformers + torch are required for backend='transformers'.") from exc
-    auto_cls = (
-        getattr(transformers, "AutoModelForImageTextToText", None)
-        or getattr(transformers, "AutoModelForVision2Seq", None)
-    )
-    if auto_cls is None:
-        raise ImportError(
-            "Neither AutoModelForImageTextToText nor AutoModelForVision2Seq is available in this "
-            "transformers version. Install transformers>=4.45 (which has AutoModelForImageTextToText) "
-            "for VL models."
-        )
-    processor = AutoProcessor.from_pretrained(
-        config.model_id, trust_remote_code=config.trust_remote_code
-    )
-    import os as _os  # noqa: PLC0415
-
-    use_accelerate = _os.environ.get("LEROBOT_TRANSFORMERS_DEVICE_MAP", "manual") != "manual"
-    # ``device_map='auto'`` triggers a known std::bad_alloc on the Qwen3-VL
-    # post-load dispatch path (the alloc fails in accelerate's hook setup
-    # even with TBs of host RAM). Default to manual: load on CPU with
-    # ``low_cpu_mem_usage=True``, then ``.to("cuda")``. Set
-    # ``LEROBOT_TRANSFORMERS_DEVICE_MAP=auto`` to opt back into the old path.
-    if use_accelerate:
-        model = auto_cls.from_pretrained(
-            config.model_id,
-            torch_dtype="auto",
-            device_map="auto",
-            low_cpu_mem_usage=True,
-            trust_remote_code=config.trust_remote_code,
-        )
-    else:
-        import torch as _torch  # noqa: PLC0415
-
-        model = auto_cls.from_pretrained(
-            config.model_id,
-            torch_dtype=_torch.bfloat16,
-            low_cpu_mem_usage=True,
-            trust_remote_code=config.trust_remote_code,
-        )
-        model = model.to("cuda")
-    model.eval()
-
-    def _gen(batch: Sequence[Sequence[dict[str, Any]]], max_tok: int, temp: float) -> list[str]:
-        outs: list[str] = []
-        for messages in batch:
-            text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
-            inputs = processor(text=[text], return_tensors="pt").to(model.device)
-            with torch.no_grad():
-                gen = model.generate(
-                    **inputs,
-                    max_new_tokens=max_tok,
-                    temperature=temp,
-                    do_sample=temp > 0.0,
-                )
-            decoded = processor.batch_decode(
-                gen[:, inputs["input_ids"].shape[-1] :], skip_special_tokens=True
-            )[0]
-            outs.append(decoded)
-        return outs
-
-    return _GenericTextClient(_gen, config)
-
-
-def _make_openai_client(config: VlmConfig) -> VlmClient:
-    """Backend that talks to any OpenAI-compatible server.
-
-    Compatible with ``vllm serve``, ``transformers serve``,
-    ``ktransformers serve``, and hosted endpoints. By default the server
-    is expected to be already running. Set ``auto_serve=True`` to have
-    this client spawn one (default: ``transformers serve``), wait until
-    it's ready, and tear it down on process exit.
-
-    Image blocks ``{"type":"image", "image":<PIL.Image>}`` are
-    auto-converted to ``image_url`` data-URLs. Video blocks
-    ``{"type":"video", "video":[<PIL>...]}`` are forwarded as
-    multi-frame ``video_url`` items where supported.
-    """
-    try:
-        from openai import OpenAI  # type: ignore[import-not-found]
-    except ImportError as exc:
-        raise ImportError(
-            "openai package is required for backend='openai'. "
-            "Install with `pip install openai`."
-        ) from exc
-
-    api_base = config.api_base
-    api_key = config.api_key
-    auto_serve = config.auto_serve
-    api_bases: list[str] = [api_base]
-
-    print(
-        f"[lerobot-annotate] backend=openai model={config.model_id} "
-        f"api_base={api_base} auto_serve={auto_serve}",
-        flush=True,
-    )
-    if auto_serve:
-        if config.parallel_servers > 1:
-            print(
-                f"[lerobot-annotate] spawning {config.parallel_servers} parallel servers",
-                flush=True,
-            )
-            api_bases = _spawn_parallel_inference_servers(config)
-        elif _server_is_up(api_base):
-            print(f"[lerobot-annotate] reusing server already up at {api_base}", flush=True)
-        else:
-            print("[lerobot-annotate] no server reachable; spawning one", flush=True)
-            api_base = _spawn_inference_server(config)
-            api_bases = [api_base]
-            print(f"[lerobot-annotate] server ready at {api_base}", flush=True)
-
-    clients = [OpenAI(base_url=base, api_key=api_key) for base in api_bases]
-    client = clients[0]
-    # round-robin counter for parallel mode
-    rr_counter = {"i": 0}
-
-    # ``mm_processor_kwargs`` is a vllm-specific extra; transformers serve
-    # rejects it with HTTP 422. Send it only when explicitly opted in via
-    # an env var (e.g. ``LEROBOT_OPENAI_SEND_MM_KWARGS=1`` for vllm).
-    send_mm_kwargs = os.environ.get(
-        "LEROBOT_OPENAI_SEND_MM_KWARGS", ""
-    ).lower() in {"1", "true", "yes"}
-
-    rr_lock = threading.Lock()
-
-    def _one_call(
-        messages: Sequence[dict[str, Any]], max_tok: int, temp: float
-    ) -> str:
-        api_messages, mm_kwargs = _to_openai_messages(messages)
-        kwargs: dict[str, Any] = {
-            "model": config.model_id,
-            "messages": api_messages,
-            "max_tokens": max_tok,
-            "temperature": temp,
-        }
-        extra_body: dict[str, Any] = {}
-        if send_mm_kwargs and mm_kwargs:
-            extra_body["mm_processor_kwargs"] = {**mm_kwargs, "do_sample_frames": True}
-        if config.chat_template_kwargs:
-            extra_body["chat_template_kwargs"] = config.chat_template_kwargs
-        if extra_body:
-            kwargs["extra_body"] = extra_body
-        with rr_lock:
-            chosen = clients[rr_counter["i"] % len(clients)]
-            rr_counter["i"] += 1
-        response = chosen.chat.completions.create(**kwargs)
-        return response.choices[0].message.content or ""
-
-    def _gen(
-        batch: Sequence[Sequence[dict[str, Any]]], max_tok: int, temp: float
-    ) -> list[str]:
-        if len(batch) <= 1 or config.client_concurrency <= 1:
-            return [_one_call(messages, max_tok, temp) for messages in batch]
-        # Parallel fan-out — vllm batches these on the server side.
-        from concurrent.futures import ThreadPoolExecutor  # noqa: PLC0415
-
-        max_workers = min(config.client_concurrency, len(batch))
-        with ThreadPoolExecutor(max_workers=max_workers) as pool:
-            futures = [
-                pool.submit(_one_call, messages, max_tok, temp) for messages in batch
-            ]
-            return [f.result() for f in futures]
-
-    return _GenericTextClient(_gen, config)
-
-
-def _spawn_parallel_inference_servers(config: VlmConfig) -> list[str]:
-    """Spawn ``config.parallel_servers`` independent vllm replicas.
-
-    Each replica:
-    - is pinned to a single GPU via ``CUDA_VISIBLE_DEVICES``
-    - listens on ``serve_port + i``
-    - is shut down via the same atexit hook as the single-server path
-
-    Returns the list of ``api_base`` URLs the client should round-robin
-    across.
-    """
-    import atexit  # noqa: PLC0415
-    import os as _os  # noqa: PLC0415
-    import shlex  # noqa: PLC0415
-    import signal  # noqa: PLC0415
-    import subprocess  # noqa: PLC0415
-    import sys  # noqa: PLC0415
-    import threading  # noqa: PLC0415
-    import time  # noqa: PLC0415
-
-    n = config.parallel_servers
-    api_bases: list[str] = []
-    procs: list[subprocess.Popen] = []
-    ready_events: list[threading.Event] = []
-    # Multiple readiness signals — uvicorn's own banner is suppressed at
-    # ``--uvicorn-log-level warning``, so we also accept vllm's own
-    # "Starting vLLM API server" line and the route-listing line. The
-    # HTTP probe below is the ultimate fallback.
-    ready_markers = (
-        "Uvicorn running",
-        "Application startup complete",
-        "Starting vLLM API server",
-        "Available routes are",
-    )
-    # Single lock for all server-stream threads so multibyte chars from
-    # different servers don't interleave and tear UTF-8 sequences.
-    print_lock = threading.Lock()
-
-    base_cmd = config.serve_command or (
-        f"vllm serve {shlex.quote(config.model_id)} "
-        f"--tensor-parallel-size 1 "
-        f"--max-model-len {config.max_model_len or 32768} "
-        f"--uvicorn-log-level warning"
-    )
-
-    num_gpus = config.num_gpus if config.num_gpus > 0 else n
-    for i in range(n):
-        port = config.serve_port + i
-        gpu = i % num_gpus
-        env = _os.environ.copy()
-        env["CUDA_VISIBLE_DEVICES"] = str(gpu)
-        cmd = base_cmd
-        if "{port}" in cmd:
-            cmd = cmd.replace("{port}", str(port))
-        else:
-            cmd = f"{cmd} --port {port}"
-        api_base = f"http://localhost:{port}/v1"
-        api_bases.append(api_base)
-        print(f"[server-{i}] launching on GPU {gpu} port {port}: {cmd}", flush=True)
-        proc = subprocess.Popen(
-            shlex.split(cmd),
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            text=True,
-            bufsize=1,
-            env=env,
-        )
-        procs.append(proc)
-        ready = threading.Event()
-        ready_events.append(ready)
-
-        def _stream(idx: int, p: subprocess.Popen, ev: threading.Event) -> None:
-            # Read whole lines and emit each line atomically under the
-            # shared print_lock so output from N servers stays readable.
-            assert p.stdout is not None
-            for line in iter(p.stdout.readline, ""):
-                with print_lock:
-                    sys.stdout.write(f"[server-{idx}] {line}")
-                    if not line.endswith(("\n", "\r")):
-                        sys.stdout.write("\n")
-                    sys.stdout.flush()
-                if any(m in line for m in ready_markers):
-                    ev.set()
-
-        threading.Thread(target=_stream, args=(i, proc, ready), daemon=True).start()
-
-        def _probe(idx: int, base: str, ev: threading.Event, p: subprocess.Popen) -> None:
-            while not ev.is_set() and p.poll() is None:
-                if _server_is_up(base):
-                    print(f"[server-{idx}] ready (http probe)", flush=True)
-                    ev.set()
-                    return
-                time.sleep(2)
-
-        threading.Thread(target=_probe, args=(i, api_base, ready, proc), daemon=True).start()
-
-    def _shutdown() -> None:
-        for i, p in enumerate(procs):
-            if p.poll() is None:
-                print(f"[server-{i}] stopping pid={p.pid}", flush=True)
-                p.send_signal(signal.SIGINT)
-        for p in procs:
-            try:
-                p.wait(timeout=15)
-            except subprocess.TimeoutExpired:
-                p.kill()
-                p.wait(timeout=5)
-
-    atexit.register(_shutdown)
-
-    deadline = time.monotonic() + config.serve_ready_timeout_s
-    while any(not ev.is_set() for ev in ready_events) and time.monotonic() < deadline:
-        for i, p in enumerate(procs):
-            if p.poll() is not None:
-                raise RuntimeError(
-                    f"[server-{i}] inference server exited unexpectedly with rc={p.returncode}"
-                )
-        time.sleep(2)
-    if any(not ev.is_set() for ev in ready_events):
-        raise RuntimeError(
-            f"[server] not all replicas became ready within {config.serve_ready_timeout_s}s"
-        )
-    print(f"[lerobot-annotate] all {n} servers ready: {api_bases}", flush=True)
-    return api_bases
-
-
-def _server_is_up(api_base: str) -> bool:
-    """Return True if ``api_base/models`` answers 200 within 2 seconds."""
-    import urllib.request  # noqa: PLC0415
-
-    url = api_base.rstrip("/") + "/models"
-    try:
-        with urllib.request.urlopen(url, timeout=2) as resp:
-            return resp.status == 200
-    except Exception:  # noqa: BLE001
-        return False
-
-
-def _spawn_inference_server(config: VlmConfig) -> str:
-    """Spawn ``transformers serve`` (or ``serve_command``), wait until it
-    accepts ``/v1/models``, and register a shutdown hook.
-
-    Streams the server's stdout/stderr to the parent terminal in
-    real-time on a background thread so users can see model-load
-    progress and errors as they happen.
-
-    Returns the full ``api_base`` URL the OpenAI client should use.
-    """
-    import atexit  # noqa: PLC0415
-    import shlex  # noqa: PLC0415
-    import signal  # noqa: PLC0415
-    import subprocess  # noqa: PLC0415
-    import sys  # noqa: PLC0415
-    import threading  # noqa: PLC0415
-    import time  # noqa: PLC0415
-    import urllib.request  # noqa: PLC0415
-
-    cmd = config.serve_command
-    if not cmd:
-        cmd = (
-            f"transformers serve {shlex.quote(config.model_id)} "
-            f"--port {config.serve_port} --continuous-batching"
-        )
-    api_base = f"http://localhost:{config.serve_port}/v1"
-    print(f"[server] launching: {cmd}", flush=True)
-    proc = subprocess.Popen(
-        shlex.split(cmd),
-        stdout=subprocess.PIPE,
-        stderr=subprocess.STDOUT,
-        text=True,
-        bufsize=1,
-    )
-
-    # Watch the server output for the uvicorn readiness banner. This is
-    # more reliable than polling /v1/models because transformers serve
-    # rescans its cache on every model-list request, which can exceed
-    # the urllib timeout and trigger an infinite probe loop.
-    ready_event = threading.Event()
-    # See _spawn_parallel_inference_servers for why we accept these.
-    ready_markers = (
-        "Uvicorn running",
-        "Application startup complete",
-        "Starting vLLM API server",
-        "Available routes are",
-    )
-
-    def _probe() -> None:
-        while not ready_event.is_set() and proc.poll() is None:
-            if _server_is_up(api_base):
-                print("[server] ready (http probe)", flush=True)
-                ready_event.set()
-                return
-            time.sleep(2)
-
-    threading.Thread(target=_probe, daemon=True).start()
-
-    def _stream_output() -> None:
-        # Read raw chunks instead of iterating lines so tqdm progress
-        # bars (which overwrite using \r) flush in real time.
-        assert proc.stdout is not None
-        buf = ""
-        prefix_started = False
-        while True:
-            ch = proc.stdout.read(1)
-            if ch == "":
-                # process exited; flush any tail
-                if buf:
-                    sys.stdout.write(buf)
-                    sys.stdout.flush()
-                return
-            if not prefix_started:
-                sys.stdout.write("[server] ")
-                prefix_started = True
-            sys.stdout.write(ch)
-            sys.stdout.flush()
-            buf += ch
-            if ch in ("\n", "\r"):
-                if any(marker in buf for marker in ready_markers):
-                    ready_event.set()
-                buf = ""
-                prefix_started = False
-
-    threading.Thread(target=_stream_output, daemon=True).start()
-
-    def _shutdown() -> None:
-        if proc.poll() is None:
-            print(f"[server] stopping pid={proc.pid}", flush=True)
-            proc.send_signal(signal.SIGINT)
-            try:
-                proc.wait(timeout=15)
-            except subprocess.TimeoutExpired:
-                proc.kill()
-                proc.wait(timeout=5)
-
-    atexit.register(_shutdown)
-
-    deadline = time.monotonic() + config.serve_ready_timeout_s
-    while time.monotonic() < deadline:
-        if proc.poll() is not None:
-            raise RuntimeError(
-                f"[server] inference server exited unexpectedly with rc={proc.returncode}. "
-                f"See [server] log lines above for the cause."
-            )
-        if ready_event.wait(timeout=2):
-            return api_base
-    proc.terminate()
-    raise RuntimeError(
-        f"[server] did not become ready within {config.serve_ready_timeout_s}s"
-    )
-
-
-def _to_openai_messages(
-    messages: Sequence[dict[str, Any]],
-) -> tuple[list[dict[str, Any]], dict[str, Any]]:
-    """Convert internal messages to OpenAI chat format.
-
-    Returns ``(api_messages, mm_kwargs)``. Multimodal-processor kwargs
-    (``fps`` from ``video_url`` blocks) are extracted out so the caller
-    can pass them via ``extra_body.mm_processor_kwargs`` rather than
-    inside the content blocks (which transformers serve rejects).
-
-    File-URL video blocks are inlined as base64 data URLs.
-    """
-    out_messages: list[dict[str, Any]] = []
-    mm_kwargs: dict[str, Any] = {}
-    for message in messages:
-        content = message.get("content")
-        if not isinstance(content, list):
-            out_messages.append({"role": message["role"], "content": content})
-            continue
-        out_blocks: list[dict[str, Any]] = []
-        for block in content:
-            block_type = block.get("type") if isinstance(block, dict) else None
-            if block_type == "text":
-                out_blocks.append({"type": "text", "text": block.get("text", "")})
-            elif block_type == "image":
-                out_blocks.append(
-                    {"type": "image_url", "image_url": {"url": _pil_to_data_url(block["image"])}}
-                )
-            elif block_type == "video":
-                frames = block.get("video", [])
-                for img in frames:
-                    out_blocks.append(
-                        {"type": "image_url", "image_url": {"url": _pil_to_data_url(img)}}
-                    )
-            elif block_type == "video_url":
-                video_url = dict(block["video_url"])
-                url = video_url.get("url", "")
-                if url.startswith("file://"):
-                    video_url["url"] = _file_to_data_url(url[len("file://") :])
-                out_blocks.append({"type": "video_url", "video_url": video_url})
-                fps = block.get("fps")
-                if fps is not None:
-                    mm_kwargs["fps"] = fps
-            else:
-                out_blocks.append(block)
-        out_messages.append({"role": message["role"], "content": out_blocks})
-    return out_messages, mm_kwargs
-
-
-def _file_to_data_url(path: str) -> str:
-    """Read a local video file and return a base64 ``data:video/mp4`` URL."""
-    import base64  # noqa: PLC0415
-
-    with open(path, "rb") as f:
-        b64 = base64.b64encode(f.read()).decode("ascii")
-    return f"data:video/mp4;base64,{b64}"
-
-
-def _pil_to_data_url(image: Any) -> str:
-    """Encode a PIL.Image as a base64 data URL."""
-    import base64  # noqa: PLC0415
-    import io  # noqa: PLC0415
-
-    buf = io.BytesIO()
-    image.save(buf, format="PNG")
-    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
-    return f"data:image/png;base64,{b64}"
-
-
-def _messages_to_prompt(messages: Sequence[dict[str, Any]]) -> Any:
-    """Pass-through hook used by the vllm backend.
-
-    vllm exposes its own multimodal entry points that vary by version; for the
-    base flow we simply forward the raw message list and let the caller's
-    custom backend handle templating. Real deployments override this.
-    """
-    return list(messages)

src/lerobot/annotations/steerable_pipeline/writer.py

@@ -1,341 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Final parquet rewrite.
-
-For every episode the writer:
-
-1. reads the staged module outputs,
-2. partitions them into a persistent slice (PERSISTENT_STYLES) and an event
-   slice (EVENT_ONLY_STYLES + style=None tool-call atoms),
-3. sorts each slice deterministically,
-4. broadcasts the persistent slice across every frame in the episode,
-5. for each frame, materializes the sublist of event rows whose timestamp
-   exactly equals that frame's timestamp,
-6. drops the legacy ``subtask_index`` column,
-7. writes the parquet shard back in place.
-
-The writer does NOT add a dataset-level ``tools`` column. Tool *calls* are
-emitted per-row via the existing ``tool_calls`` field on the v3.1 row
-struct (PR 1) for every speech atom. The tool *schema* (the description
-of the ``say`` function and its parameters) is a fixed code constant —
-``SAY_TOOL_SCHEMA`` below — and downstream chat-template consumers import
-it directly rather than reading a redundant per-row column.
-
-Invariants enforced here (and re-checked by the validator):
-
-- per-episode persistent slice is byte-identical across every frame;
-- ``language_events`` rows on a frame all have ``timestamp == frame_ts``
-  (timestamps come straight from the source parquet — never recomputed);
-- every row passes ``column_for_style(style)``.
-"""
-
-from __future__ import annotations
-
-import logging
-from collections import defaultdict
-from collections.abc import Iterable, Sequence
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import pyarrow as pa
-import pyarrow.parquet as pq
-
-from lerobot.datasets.language import (
-    EVENT_ONLY_STYLES,
-    LANGUAGE_EVENTS,
-    LANGUAGE_PERSISTENT,
-    PERSISTENT_STYLES,
-    column_for_style,
-    validate_camera_field,
-)
-
-from .reader import EpisodeRecord
-from .staging import EpisodeStaging
-
-logger = logging.getLogger(__name__)
-
-
-# Tool schema constants moved to lerobot.datasets.language in PR 1 — single
-# source of truth. Re-exported here so existing imports
-# (``from lerobot.annotations.steerable_pipeline.writer import SAY_TOOL_SCHEMA``)
-# keep working.
-from lerobot.datasets.language import DEFAULT_TOOLS, SAY_TOOL_SCHEMA  # noqa: F401, E402
-
-
-def _row_persistent_sort_key(row: dict[str, Any]) -> tuple:
-    return (float(row["timestamp"]), row.get("style") or "", row.get("role") or "")
-
-
-def _row_event_sort_key(row: dict[str, Any]) -> tuple:
-    # events are bucketed per-frame, but within a frame we still want determinism
-    return (
-        row.get("style") or "",
-        row.get("role") or "",
-        row.get("camera") or "",
-    )
-
-
-def _normalize_persistent_row(row: dict[str, Any]) -> dict[str, Any]:
-    """Coerce a staged row into the persistent column's struct shape."""
-    style = row.get("style")
-    if style not in PERSISTENT_STYLES:
-        raise ValueError(
-            f"persistent slice contains row with non-persistent style {style!r}; "
-            "row would be misrouted under column_for_style()"
-        )
-    if "timestamp" not in row:
-        raise ValueError(f"persistent row missing timestamp: {row!r}")
-    camera = row.get("camera")
-    validate_camera_field(style, camera)
-    return {
-        "role": str(row["role"]),
-        "content": None if row.get("content") is None else str(row["content"]),
-        "style": style,
-        "timestamp": float(row["timestamp"]),
-        "camera": None if camera is None else str(camera),
-        "tool_calls": _normalize_tool_calls(row.get("tool_calls")),
-    }
-
-
-def _normalize_event_row(row: dict[str, Any]) -> dict[str, Any]:
-    """Coerce a staged row into the event column's struct shape (no timestamp)."""
-    style = row.get("style")
-    if style is not None and style not in EVENT_ONLY_STYLES:
-        raise ValueError(
-            f"event slice contains row with style {style!r}; expected None or one of {EVENT_ONLY_STYLES}"
-        )
-    if column_for_style(style) != LANGUAGE_EVENTS:
-        raise ValueError(f"event row with style {style!r} would not route to language_events")
-    camera = row.get("camera")
-    validate_camera_field(style, camera)
-    return {
-        "role": str(row["role"]),
-        "content": None if row.get("content") is None else str(row["content"]),
-        "style": style,
-        "camera": None if camera is None else str(camera),
-        "tool_calls": _normalize_tool_calls(row.get("tool_calls")),
-    }
-
-
-def _normalize_tool_calls(value: Any) -> list[Any] | None:
-    if value is None:
-        return None
-    if not isinstance(value, list):
-        raise ValueError(f"tool_calls must be a list or None, got {type(value).__name__}")
-    return list(value)
-
-
-def _validate_atom_invariants(row: dict[str, Any]) -> None:
-    """At-least-one of content/tool_calls; style=None implies tool_calls."""
-    has_content = row.get("content") is not None
-    has_tools = row.get("tool_calls") is not None
-    if not (has_content or has_tools):
-        raise ValueError(f"row has neither content nor tool_calls: {row!r}")
-    if row.get("style") is None and not has_tools:
-        raise ValueError(f"style=None requires tool_calls: {row!r}")
-
-
-def _validate_speech_atom(row: dict[str, Any]) -> None:
-    """Speech atoms: role=assistant, style=None, content=None, say tool call."""
-    if row.get("style") is not None:
-        return  # not a speech atom
-    if row.get("role") != "assistant":
-        raise ValueError(f"speech atom must have role=assistant: {row!r}")
-    if row.get("content") is not None:
-        raise ValueError(f"speech atom must have content=null: {row!r}")
-    tool_calls = row.get("tool_calls")
-    if not tool_calls or not isinstance(tool_calls, list):
-        raise ValueError(f"speech atom must have non-empty tool_calls list: {row!r}")
-    first = tool_calls[0]
-    if not isinstance(first, dict):
-        raise ValueError(f"speech atom tool_calls[0] must be a dict: {row!r}")
-    if first.get("type") != "function":
-        raise ValueError(f"speech atom tool_calls[0].type must be 'function': {row!r}")
-    fn = first.get("function") or {}
-    if fn.get("name") != "say":
-        raise ValueError(f"speech atom tool_calls[0].function.name must be 'say': {row!r}")
-    args = fn.get("arguments") or {}
-    if not isinstance(args, dict) or "text" not in args or not isinstance(args["text"], str):
-        raise ValueError(f"speech atom must carry 'text' string in arguments: {row!r}")
-
-
-@dataclass
-class LanguageColumnsWriter:
-    """Rewrite ``data/chunk-*/file-*.parquet`` with the two language columns."""
-
-    drop_existing_subtask_index: bool = True
-
-    def write_all(
-        self,
-        records: Sequence[EpisodeRecord],
-        staging_dir: Path,
-        root: Path,
-    ) -> list[Path]:
-        episodes_by_path: dict[Path, list[EpisodeRecord]] = defaultdict(list)
-        for record in records:
-            episodes_by_path[record.data_path].append(record)
-
-        written: list[Path] = []
-        for path, eps in episodes_by_path.items():
-            self._rewrite_one(path, eps, staging_dir, root)
-            written.append(path)
-        return written
-
-    def _rewrite_one(
-        self,
-        path: Path,
-        episodes: Sequence[EpisodeRecord],
-        staging_dir: Path,
-        root: Path,
-    ) -> None:
-        table = pq.read_table(path)
-        n_rows = table.num_rows
-
-        # Ensure we cover every episode in the file. Episodes that don't have
-        # staging artifacts are passed through with empty annotation lists —
-        # this keeps the writer idempotent and safe for partial reruns.
-        staged_per_ep: dict[int, dict[str, list[dict[str, Any]]]] = {}
-        for record in episodes:
-            staging = EpisodeStaging(staging_dir, record.episode_index)
-            staged_per_ep[record.episode_index] = staging.read_all()
-
-        persistent_by_ep: dict[int, list[dict[str, Any]]] = {}
-        events_by_ep_ts: dict[int, dict[float, list[dict[str, Any]]]] = {}
-
-        for ep_index, ep_staged in staged_per_ep.items():
-            persistent_rows: list[dict[str, Any]] = []
-            event_rows: list[dict[str, Any]] = []  # carry timestamp until bucketed
-            for _module_name, rows in ep_staged.items():
-                for row in rows:
-                    style = row.get("style")
-                    if column_for_style(style) == LANGUAGE_PERSISTENT:
-                        persistent_rows.append(row)
-                    else:
-                        event_rows.append(row)
-
-            persistent_rows.sort(key=_row_persistent_sort_key)
-            normalized_persistent = []
-            for r in persistent_rows:
-                _validate_atom_invariants(r)
-                _validate_speech_atom(r)
-                normalized_persistent.append(_normalize_persistent_row(r))
-            persistent_by_ep[ep_index] = normalized_persistent
-
-            buckets: dict[float, list[dict[str, Any]]] = defaultdict(list)
-            for r in event_rows:
-                _validate_atom_invariants(r)
-                _validate_speech_atom(r)
-                ts = float(r["timestamp"])
-                buckets[ts].append(_normalize_event_row(r))
-            for ts in list(buckets.keys()):
-                buckets[ts].sort(key=_row_event_sort_key)
-            events_by_ep_ts[ep_index] = buckets
-
-        episode_col = (
-            table.column("episode_index").to_pylist() if "episode_index" in table.column_names else None
-        )
-        ts_col = table.column("timestamp").to_pylist() if "timestamp" in table.column_names else None
-        if episode_col is None or ts_col is None:
-            raise ValueError(f"{path} is missing 'episode_index' or 'timestamp' — required by the writer.")
-
-        per_row_persistent: list[list[dict[str, Any]]] = []
-        per_row_events: list[list[dict[str, Any]]] = []
-        for i in range(n_rows):
-            ep = episode_col[i]
-            ts = float(ts_col[i])
-            per_row_persistent.append(persistent_by_ep.get(ep, []))
-            buckets = events_by_ep_ts.get(ep, {})
-            per_row_events.append(buckets.get(ts, []))
-
-        new_table = self._materialize_table(
-            table, per_row_persistent, per_row_events, drop_old=self.drop_existing_subtask_index
-        )
-        pq.write_table(new_table, path)
-
-    def _materialize_table(
-        self,
-        table: pa.Table,
-        persistent: list[list[dict[str, Any]]],
-        events: list[list[dict[str, Any]]],
-        *,
-        drop_old: bool,
-    ) -> pa.Table:
-        cols = []
-        names = []
-        for name in table.column_names:
-            if drop_old and name == "subtask_index":
-                continue
-            if name in (LANGUAGE_PERSISTENT, LANGUAGE_EVENTS):
-                continue  # we'll re-add canonical versions
-            # Strip any legacy ``tools`` column previously emitted by older
-            # writers — the schema no longer uses it (constant lives in
-            # SAY_TOOL_SCHEMA / DEFAULT_TOOLS).
-            if name == "tools":
-                continue
-            cols.append(table.column(name))
-            names.append(name)
-
-        # We let pyarrow infer struct/list schema rather than passing the
-        # canonical type from `lerobot.datasets.language` directly: that type
-        # uses `pa.json_()` for the `tool_calls` element type, which
-        # `pa.array(..., type=...)` cannot materialize from Python lists on
-        # current pyarrow versions. The inferred schema round-trips through
-        # parquet and `LeRobotDataset` correctly — see PR 1's
-        # `tests/datasets/test_language.py` which exercises the same flow.
-        persistent_arr = pa.array(persistent)
-        events_arr = pa.array(events)
-
-        cols.extend([persistent_arr, events_arr])
-        names.extend([LANGUAGE_PERSISTENT, LANGUAGE_EVENTS])
-
-        return pa.Table.from_arrays(cols, names=names)
-
-
-def speech_atom(timestamp: float, text: str) -> dict[str, Any]:
-    """Build a canonical speech tool-call atom for the events column."""
-    return {
-        "role": "assistant",
-        "content": None,
-        "style": None,
-        "timestamp": float(timestamp),
-        "camera": None,
-        "tool_calls": [
-            {
-                "type": "function",
-                "function": {
-                    "name": "say",
-                    "arguments": {"text": text},
-                },
-            }
-        ],
-    }
-
-
-def normalize_rows_for_writer(
-    rows: Iterable[dict[str, Any]],
-) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
-    """Helper used by tests/validators to partition a flat row list into
-    (persistent_rows, event_rows) using ``column_for_style``.
-    """
-    persistent: list[dict[str, Any]] = []
-    events: list[dict[str, Any]] = []
-    for row in rows:
-        if column_for_style(row.get("style")) == LANGUAGE_PERSISTENT:
-            persistent.append(row)
-        else:
-            events.append(row)
-    return persistent, events

src/lerobot/configs/__init__.py

@@ -21,9 +21,9 @@ are intentionally NOT re-exported here to avoid circular dependencies
 Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
 """
 
+from .dataset import DatasetRecordConfig
 from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
-from .recipe import MessageTurn, TrainingRecipe, load_recipe
 from .types import (
     FeatureType,
     NormalizationMode,
@@ -40,12 +40,10 @@ __all__ = [
     "PolicyFeature",
     "RTCAttentionSchedule",
     # Config classes
+    "DatasetRecordConfig",
     "DatasetConfig",
     "EvalConfig",
-    "MessageTurn",
     "PeftConfig",
     "PreTrainedConfig",
-    "TrainingRecipe",
     "WandBConfig",
-    "load_recipe",
 ]

src/lerobot/configs/dataset.py

@@ -0,0 +1,77 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared dataset recording configuration used by both ``lerobot-record`` and ``lerobot-rollout``."""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+
+@dataclass
+class DatasetRecordConfig:
+    # Dataset identifier. By convention it should match '{hf_username}/{dataset_name}' (e.g. `lerobot/test`).
+    repo_id: str = ""
+    # A short but accurate description of the task performed during the recording (e.g. "Pick the Lego block and drop it in the box on the right.")
+    single_task: str = ""
+    # Root directory where the dataset will be stored (e.g. 'dataset/path'). If None, defaults to $HF_LEROBOT_HOME/repo_id.
+    root: str | Path | None = None
+    # Limit the frames per second.
+    fps: int = 30
+    # Number of seconds for data recording for each episode.
+    episode_time_s: int | float = 60
+    # Number of seconds for resetting the environment after each episode.
+    reset_time_s: int | float = 60
+    # Number of episodes to record.
+    num_episodes: int = 50
+    # Encode frames in the dataset into video
+    video: bool = True
+    # Upload dataset to Hugging Face hub.
+    push_to_hub: bool = True
+    # Upload on private repository on the Hugging Face hub.
+    private: bool = False
+    # Add tags to your dataset on the hub.
+    tags: list[str] | None = None
+    # Number of subprocesses handling the saving of frames as PNG. Set to 0 to use threads only;
+    # set to ≥1 to use subprocesses, each using threads to write images. The best number of processes
+    # and threads depends on your system. We recommend 4 threads per camera with 0 processes.
+    # If fps is unstable, adjust the thread count. If still unstable, try using 1 or more subprocesses.
+    num_image_writer_processes: int = 0
+    # Number of threads writing the frames as png images on disk, per camera.
+    # Too many threads might cause unstable teleoperation fps due to main thread being blocked.
+    # Not enough threads might cause low camera fps.
+    num_image_writer_threads_per_camera: int = 4
+    # Number of episodes to record before batch encoding videos
+    # Set to 1 for immediate encoding (default behavior), or higher for batched encoding
+    video_encoding_batch_size: int = 1
+    # Video codec for encoding videos. Options: 'h264', 'hevc', 'libsvtav1', 'auto',
+    # or hardware-specific: 'h264_videotoolbox', 'h264_nvenc', 'h264_vaapi', 'h264_qsv'.
+    # Use 'auto' to auto-detect the best available hardware encoder.
+    vcodec: str = "libsvtav1"
+    # Enable streaming video encoding: encode frames in real-time during capture instead
+    # of writing PNG images first. Makes save_episode() near-instant. More info in the documentation: https://huggingface.co/docs/lerobot/streaming_video_encoding
+    streaming_encoding: bool = False
+    # Maximum number of frames to buffer per camera when using streaming encoding.
+    # ~1s buffer at 30fps. Provides backpressure if the encoder can't keep up.
+    encoder_queue_maxsize: int = 30
+    # Number of threads per encoder instance. None = auto (codec default).
+    # Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
+    encoder_threads: int | None = None
+    # Rename map for the observation to override the image and state keys
+    rename_map: dict[str, str] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if self.repo_id:
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            self.repo_id = f"{self.repo_id}_{timestamp}"

src/lerobot/configs/recipe.py

@@ -1,193 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import re
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any, Literal, get_args
-
-MessageRole = Literal["user", "assistant", "system", "tool"]
-MessageStream = Literal["high_level", "low_level"]
-
-DEFAULT_BINDINGS = {
-    "subtask": "active_at(t, style=subtask)",
-    "memory": "active_at(t, style=memory)",
-    "plan": "active_at(t, style=plan)",
-    "speech": "emitted_at(t, role=assistant, tool_name=say)",
-    "interjection": "emitted_at(t, style=interjection)",
-    "vqa": "emitted_at(t, style=vqa, role=assistant)",
-    "vqa_query": "emitted_at(t, style=vqa, role=user)",
-}
-
-_PLACEHOLDER_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")
-_VALID_ROLES = frozenset(get_args(MessageRole))
-_VALID_STREAMS = frozenset(get_args(MessageStream))
-
-
-@dataclass
-class MessageTurn:
-    """A single chat-style turn in a recipe template.
-
-    ``content`` may be a plain string, a list of HF-style multimodal blocks, or
-    ``None`` when ``tool_calls_from`` supplies tool-call payloads instead.
-    ``stream`` tags the turn for downstream filtering, ``target`` flags it as a
-    training target, and ``if_present`` skips the turn when the named binding
-    resolves to ``None``.
-    """
-
-    role: MessageRole
-    content: str | list[dict[str, Any]] | None = None
-    stream: MessageStream | None = None
-    target: bool = False
-    if_present: str | None = None
-    tool_calls_from: str | None = None
-
-    def __post_init__(self) -> None:
-        """Validate role, stream, and content after dataclass construction."""
-        if self.role not in _VALID_ROLES:
-            raise ValueError(f"Unsupported message role: {self.role!r}")
-        if self.stream is not None and self.stream not in _VALID_STREAMS:
-            raise ValueError(f"Unsupported message stream: {self.stream!r}")
-        if self.content is None and self.tool_calls_from is None:
-            raise ValueError("MessageTurn.content is required unless tool_calls_from is set.")
-        if self.content is not None and not isinstance(self.content, (str, list)):
-            raise TypeError("MessageTurn.content must be a string, a list of HF-style blocks, or None.")
-        if isinstance(self.content, list):
-            for block in self.content:
-                if not isinstance(block, dict) or "type" not in block:
-                    raise ValueError(
-                        "Multimodal content blocks must be HF-style dictionaries with a type key."
-                    )
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> MessageTurn:
-        """Construct a :class:`MessageTurn` from a plain dictionary."""
-        return cls(**data)
-
-
-@dataclass
-class TrainingRecipe:
-    """A recipe describing how to render training samples from language rows.
-
-    A recipe is either a *message recipe* (``messages`` plus optional
-    ``bindings``) or a *blend recipe* (``blend`` mapping names to weighted
-    sub-recipes). ``weight`` is only meaningful inside a blend.
-    """
-
-    messages: list[MessageTurn] | None = None
-    bindings: dict[str, str] | None = None
-    blend: dict[str, TrainingRecipe] | None = None
-    weight: float | None = None
-
-    def __post_init__(self) -> None:
-        """Validate that exactly one of ``messages`` or ``blend`` is set."""
-        if self.messages is not None and self.blend is not None:
-            raise ValueError("TrainingRecipe must set only one of messages or blend.")
-        if self.messages is None and self.blend is None:
-            raise ValueError("TrainingRecipe must set one of messages or blend.")
-
-        if self.messages is not None:
-            self._validate_message_recipe()
-        if self.blend is not None:
-            self._validate_blend_recipe()
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any]) -> TrainingRecipe:
-        """Construct a :class:`TrainingRecipe` from a nested dictionary."""
-        data = dict(data)
-        if data.get("messages") is not None:
-            data["messages"] = [
-                turn if isinstance(turn, MessageTurn) else MessageTurn.from_dict(turn)
-                for turn in data["messages"]
-            ]
-        if data.get("blend") is not None:
-            data["blend"] = {
-                name: recipe if isinstance(recipe, TrainingRecipe) else cls.from_dict(recipe)
-                for name, recipe in data["blend"].items()
-            }
-        return cls(**data)
-
-    @classmethod
-    def from_yaml(cls, path: str | Path) -> TrainingRecipe:
-        """Load a :class:`TrainingRecipe` from a YAML file at ``path``."""
-        import yaml  # type: ignore[import-untyped]
-
-        with open(path) as f:
-            data = yaml.safe_load(f)
-        if not isinstance(data, dict):
-            raise ValueError(f"Recipe YAML must contain a mapping at the top level: {path}")
-        return cls.from_dict(data)
-
-    def _validate_message_recipe(self) -> None:
-        """Ensure every templated binding is known and at least one turn is a target."""
-        assert self.messages is not None
-        known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"}
-
-        for turn in self.messages:
-            missing = self._referenced_bindings(turn) - known_bindings
-            if missing:
-                raise ValueError(f"MessageTurn references unknown binding(s): {sorted(missing)}")
-
-        if not any(turn.target for turn in self.messages):
-            raise ValueError("Message recipes must contain at least one target turn.")
-
-    def _validate_blend_recipe(self) -> None:
-        """Ensure each blend component is a non-empty, weighted message recipe."""
-        assert self.blend is not None
-        if not self.blend:
-            raise ValueError("Blend recipes must contain at least one component.")
-
-        for name, recipe in self.blend.items():
-            if recipe.blend is not None:
-                raise ValueError(f"Blend component {name!r} cannot itself define a blend.")
-            if recipe.messages is None:
-                raise ValueError(f"Blend component {name!r} must define messages.")
-            if recipe.weight is None:
-                raise ValueError(f"Blend component {name!r} must define weight.")
-            if recipe.weight <= 0:
-                raise ValueError(f"Blend component {name!r} must have a positive weight.")
-
-    def _referenced_bindings(self, turn: MessageTurn) -> set[str]:
-        """Return the binding names that ``turn`` references via placeholders or attributes."""
-        names: set[str] = set()
-        if turn.if_present is not None:
-            names.add(turn.if_present)
-        if turn.tool_calls_from is not None:
-            names.add(turn.tool_calls_from)
-        names.update(_placeholders_in_content(turn.content))
-        return names
-
-
-def _placeholders_in_content(content: str | list[dict[str, Any]] | None) -> set[str]:
-    """Return the set of ``${name}`` placeholders found anywhere in ``content``."""
-    if content is None:
-        return set()
-    if isinstance(content, str):
-        return set(_PLACEHOLDER_RE.findall(content))
-
-    names: set[str] = set()
-    for block in content:
-        for value in block.values():
-            if isinstance(value, str):
-                names.update(_PLACEHOLDER_RE.findall(value))
-    return names
-
-
-def load_recipe(path: str | Path) -> TrainingRecipe:
-    """Load a :class:`TrainingRecipe` from a YAML file at ``path``."""
-    return TrainingRecipe.from_yaml(path)

src/lerobot/configs/recipes/pi05_hirobot.yaml

@@ -1,74 +0,0 @@
-blend:
-
-  memory_update:
-    weight: 0.10
-    bindings:
-      prior_memory: "nth_prev(style=memory, offset=1)"
-      current_memory: "emitted_at(t, style=memory)"
-      completed_subtask: "nth_prev(style=subtask, offset=1)"
-    messages:
-      - {role: user, content: "${task}", stream: high_level}
-      - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
-      - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
-      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
-
-  user_interjection_response:
-    weight: 0.16
-    bindings:
-      prior_plan: "nth_prev(style=plan, offset=1)"
-      current_plan: "emitted_at(t, style=plan)"
-      interjection: "emitted_at(t, style=interjection)"
-      speech: "emitted_at(t, role=assistant, tool_name=say)"
-    messages:
-      - {role: user, content: "${task}", stream: high_level}
-      - {role: assistant, content: "Previous plan:\n${prior_plan}", stream: high_level, if_present: prior_plan}
-      - {role: user, content: "${interjection}", stream: high_level, if_present: interjection}
-      - {role: assistant, content: "${current_plan}", stream: high_level, target: true, if_present: current_plan, tool_calls_from: speech}
-
-  high_level_subtask:
-    weight: 0.15
-    bindings:
-      next_subtask: "nth_next(style=subtask, offset=1)"
-    messages:
-      - {role: user, content: "${task}\nPlan: ${plan}\nMemory: ${memory}", stream: high_level}
-      - {role: user, content: "Current subtask: ${subtask}", stream: high_level, if_present: subtask}
-      - {role: assistant, content: "${next_subtask}", stream: high_level, target: true}
-
-  low_level_execution:
-    weight: 0.35
-    messages:
-      - {role: user, content: "${task}\nPlan: ${plan}\nMemory: ${memory}", stream: high_level}
-      - {role: assistant, content: "${subtask}", stream: low_level, target: true}
-
-  # VQA is view-dependent: bbox / keypoint / count answers only make sense for
-  # the camera they were grounded against. Each camera gets its own sub-recipe
-  # so the resolver can disambiguate via `camera=...` and the user-turn carries
-  # the matching image block. Adjust the camera keys (and add more sub-recipes)
-  # to match the cameras present on your dataset.
-  ask_vqa_top:
-    weight: 0.10
-    bindings:
-      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.top)"
-      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)"
-    messages:
-      - role: user
-        stream: high_level
-        if_present: vqa_query
-        content:
-          - {type: image, feature: observation.images.top}
-          - {type: text, text: "${vqa_query}"}
-      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
-
-  ask_vqa_wrist:
-    weight: 0.10
-    bindings:
-      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)"
-      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)"
-    messages:
-      - role: user
-        stream: high_level
-        if_present: vqa_query
-        content:
-          - {type: image, feature: observation.images.wrist}
-          - {type: text, text: "${vqa_query}"}
-      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}

src/lerobot/configs/recipes/smolvla2_hirobot.yaml

@@ -1,88 +0,0 @@
-# SmolVLA2 canonical training recipe — Hi Robot / MEM / ECoT blend.
-#
-# Same blend shape as pi05_hirobot.yaml. SmolVLA2 differs from Pi0.5 in
-# how the renderer's output is consumed:
-#
-#   - SmolVLA2 calls SmolVLM's tokenizer.apply_chat_template(messages,
-#     tools=DEFAULT_TOOLS) on the rendered messages, since SmolVLM is a
-#     chat-pretrained backbone.
-#   - The processor builds a `text_labels` tensor that masks every token
-#     except those belonging to messages whose index is in
-#     `target_message_indices`. Cross-entropy on those positions trains
-#     the LM head.
-#   - `predict_actions = bool(targets_by_stream.get("low_level"))` —
-#     same convention as Pi0.5. ``low_level_execution`` is the only
-#     branch that runs the action expert / flow head.
-
-blend:
-
-  memory_update:
-    weight: 0.10
-    bindings:
-      prior_memory: "nth_prev(style=memory, offset=1)"
-      current_memory: "emitted_at(t, style=memory)"
-      completed_subtask: "nth_prev(style=subtask, offset=1)"
-    messages:
-      - {role: user, content: "${task}", stream: high_level}
-      - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
-      - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
-      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
-
-  user_interjection_response:
-    weight: 0.16
-    bindings:
-      prior_plan: "nth_prev(style=plan, offset=1)"
-      current_plan: "emitted_at(t, style=plan)"
-      interjection: "emitted_at(t, style=interjection)"
-      speech: "emitted_at(t, role=assistant, tool_name=say)"
-    messages:
-      - {role: user, content: "${task}", stream: high_level}
-      - {role: assistant, content: "Previous plan:\n${prior_plan}", stream: high_level, if_present: prior_plan}
-      - {role: user, content: "${interjection}", stream: high_level, if_present: interjection}
-      - {role: assistant, content: "${current_plan}", stream: high_level, target: true, if_present: current_plan, tool_calls_from: speech}
-
-  high_level_subtask:
-    weight: 0.15
-    bindings:
-      next_subtask: "nth_next(style=subtask, offset=1)"
-    messages:
-      - {role: user, content: "${task}\nPlan: ${plan}\nMemory: ${memory}", stream: high_level}
-      - {role: user, content: "Current subtask: ${subtask}", stream: high_level, if_present: subtask}
-      - {role: assistant, content: "${next_subtask}", stream: high_level, target: true}
-
-  low_level_execution:
-    weight: 0.35
-    messages:
-      - {role: user, content: "${task}\nPlan: ${plan}\nMemory: ${memory}", stream: high_level}
-      - {role: assistant, content: "${subtask}", stream: low_level, target: true}
-
-  # Per-camera VQA sub-recipes (PR 1's view-dependent style routing).
-  # Adjust the camera keys (and add more sub-recipes) to match the
-  # cameras present on your dataset.
-  ask_vqa_top:
-    weight: 0.10
-    bindings:
-      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.top)"
-      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)"
-    messages:
-      - role: user
-        stream: high_level
-        if_present: vqa_query
-        content:
-          - {type: image, feature: observation.images.top}
-          - {type: text, text: "${vqa_query}"}
-      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
-
-  ask_vqa_wrist:
-    weight: 0.10
-    bindings:
-      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)"
-      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)"
-    messages:
-      - role: user
-        stream: high_level
-        if_present: vqa_query
-        content:
-          - {type: image, feature: observation.images.wrist}
-          - {type: text, text: "${vqa_query}"}
-      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}

src/lerobot/datasets/__init__.py

@@ -37,14 +37,6 @@ from .dataset_tools import (
 from .factory import make_dataset, resolve_delta_timestamps
 from .image_writer import safe_stop_image_writer
 from .io_utils import load_episodes, write_stats
-from .language import (
-    EVENT_ONLY_STYLES,
-    LANGUAGE_EVENTS,
-    LANGUAGE_PERSISTENT,
-    PERSISTENT_STYLES,
-    STYLE_REGISTRY,
-    column_for_style,
-)
 from .lerobot_dataset import LeRobotDataset
 from .multi_dataset import MultiLeRobotDataset
 from .pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
@@ -61,15 +53,10 @@ __all__ = [
     "CODEBASE_VERSION",
     "DEFAULT_EPISODES_PATH",
     "DEFAULT_QUANTILES",
-    "EVENT_ONLY_STYLES",
     "EpisodeAwareSampler",
-    "LANGUAGE_EVENTS",
-    "LANGUAGE_PERSISTENT",
     "LeRobotDataset",
     "LeRobotDatasetMetadata",
     "MultiLeRobotDataset",
-    "PERSISTENT_STYLES",
-    "STYLE_REGISTRY",
     "StreamingLeRobotDataset",
     "VideoEncodingManager",
     "add_features",
@@ -79,7 +66,6 @@ __all__ = [
     "convert_image_to_video_dataset",
     "create_initial_features",
     "create_lerobot_dataset_card",
-    "column_for_style",
     "delete_episodes",
     "get_feature_stats",
     "load_episodes",

src/lerobot/datasets/compute_stats.py

@@ -512,7 +512,7 @@ def compute_episode_stats(
 
     ep_stats = {}
     for key, data in episode_data.items():
-        if features[key]["dtype"] in {"string", "language"}:
+        if features[key]["dtype"] == "string":
             continue
 
         if features[key]["dtype"] in ["image", "video"]:

src/lerobot/datasets/dataset_metadata.py

@@ -34,6 +34,7 @@ from .io_utils import (
     load_episodes,
     load_info,
     load_stats,
+    load_subtasks,
     load_tasks,
     write_info,
     write_json,
@@ -176,6 +177,7 @@ class LeRobotDatasetMetadata:
         self.info = load_info(self.root)
         check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
         self.tasks = load_tasks(self.root)
+        self.subtasks = load_subtasks(self.root)
         self.episodes = load_episodes(self.root)
         self.stats = load_stats(self.root)
 
@@ -318,28 +320,6 @@ class LeRobotDatasetMetadata:
         """Keys to access visual modalities (regardless of their storage method)."""
         return [key for key, ft in self.features.items() if ft["dtype"] in ["video", "image"]]
 
-    @property
-    def tools(self) -> list[dict]:
-        """OpenAI-style tool schemas declared by this dataset.
-
-        Read from ``meta/info.json["tools"]``. Returns a copy, so callers
-        can mutate the result safely. Falls back to
-        :data:`lerobot.datasets.language.DEFAULT_TOOLS` (the canonical
-        ``say`` schema) when the dataset doesn't declare any — that way
-        unannotated datasets and chat-template consumers
-        (``apply_chat_template(messages, tools=meta.tools)``) keep
-        working out of the box.
-
-        Implementations live under :mod:`lerobot.tools` (one file per
-        tool); see ``docs/source/tools.mdx`` for the authoring guide.
-        """
-        from .language import DEFAULT_TOOLS  # noqa: PLC0415  (avoid circular import)
-
-        declared = self.info.get("tools")
-        if isinstance(declared, list) and declared:
-            return [dict(t) for t in declared]
-        return [dict(t) for t in DEFAULT_TOOLS]
-
     @property
     def names(self) -> dict[str, list | dict]:
         """Names of the various dimensions of vector modalities."""
@@ -655,6 +635,7 @@ class LeRobotDatasetMetadata:
         _validate_feature_names(features)
 
         obj.tasks = None
+        obj.subtasks = None
         obj.episodes = None
         obj.stats = None
         obj.info = create_empty_dataset_info(

src/lerobot/datasets/dataset_reader.py

@@ -295,4 +295,9 @@ class DatasetReader:
         task_idx = item["task_index"].item()
         item["task"] = self._meta.tasks.iloc[task_idx].name
 
+        # add subtask information if available
+        if "subtask_index" in self._meta.features and self._meta.subtasks is not None:
+            subtask_idx = item["subtask_index"].item()
+            item["subtask"] = self._meta.subtasks.iloc[subtask_idx].name
+
         return item

src/lerobot/datasets/feature_utils.py

@@ -22,12 +22,6 @@ from PIL import Image as PILImage
 from lerobot.utils.constants import DEFAULT_FEATURES
 from lerobot.utils.utils import is_valid_numpy_dtype_string
 
-from .language import (
-    LANGUAGE_PERSISTENT,
-    is_language_column,
-    language_events_column_feature,
-    language_persistent_column_feature,
-)
 from .utils import (
     DEFAULT_CHUNK_SIZE,
     DEFAULT_DATA_FILE_SIZE_IN_MB,
@@ -51,13 +45,7 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
     """
     hf_features = {}
     for key, ft in features.items():
-        if is_language_column(key):
-            hf_features[key] = (
-                language_persistent_column_feature()
-                if key == LANGUAGE_PERSISTENT
-                else language_events_column_feature()
-            )
-        elif ft["dtype"] == "video":
+        if ft["dtype"] == "video":
             continue
         elif ft["dtype"] == "image":
             hf_features[key] = datasets.Image()
@@ -254,8 +242,6 @@ def validate_feature_dtype_and_shape(
         return validate_feature_image_or_video(name, expected_shape, value)
     elif expected_dtype == "string":
         return validate_feature_string(name, value)
-    elif expected_dtype == "language":
-        return ""
     else:
         raise NotImplementedError(f"The feature dtype '{expected_dtype}' is not implemented yet.")
 

src/lerobot/datasets/io_utils.py

@@ -34,6 +34,7 @@ from lerobot.utils.utils import SuppressProgressBars, flatten_dict, unflatten_di
 from .utils import (
     DEFAULT_DATA_FILE_SIZE_IN_MB,
     DEFAULT_EPISODES_PATH,
+    DEFAULT_SUBTASKS_PATH,
     DEFAULT_TASKS_PATH,
     EPISODES_DIR,
     INFO_PATH,
@@ -188,6 +189,14 @@ def load_tasks(local_dir: Path) -> pandas.DataFrame:
     return tasks
 
 
+def load_subtasks(local_dir: Path) -> pandas.DataFrame | None:
+    """Load subtasks from subtasks.parquet if it exists."""
+    subtasks_path = local_dir / DEFAULT_SUBTASKS_PATH
+    if subtasks_path.exists():
+        return pd.read_parquet(subtasks_path)
+    return None
+
+
 def write_episodes(episodes: Dataset, local_dir: Path) -> None:
     """Write episode metadata to a parquet file in the LeRobot v3.0 format.
     This function writes episode-level metadata to a single parquet file.
@@ -259,13 +268,11 @@ def hf_transform_to_torch(items_dict: dict[str, list[Any]]) -> dict[str, list[to
         dict: The batch with items converted to torch tensors.
     """
     for key in items_dict:
-        if key in {"language_persistent", "language_events"}:
-            continue
         first_item = items_dict[key][0]
         if isinstance(first_item, PILImage.Image):
             to_tensor = transforms.ToTensor()
             items_dict[key] = [to_tensor(img) for img in items_dict[key]]
-        elif first_item is None or isinstance(first_item, dict):
+        elif first_item is None:
             pass
         else:
             items_dict[key] = [x if isinstance(x, str) else torch.tensor(x) for x in items_dict[key]]
@@ -301,11 +308,7 @@ def item_to_torch(item: dict) -> dict:
         dict: Dictionary with all tensor-like items converted to torch.Tensor.
     """
     for key, val in item.items():
-        if isinstance(val, (np.ndarray | list)) and key not in [
-            "task",
-            "language_persistent",
-            "language_events",
-        ]:
+        if isinstance(val, (np.ndarray | list)) and key not in ["task"]:
             # Convert numpy arrays and lists to torch tensors
             item[key] = torch.tensor(val)
     return item

src/lerobot/datasets/language.py

@@ -1,236 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from typing import Literal
-
-import datasets
-import pyarrow as pa
-
-LANGUAGE_PERSISTENT = "language_persistent"
-LANGUAGE_EVENTS = "language_events"
-LANGUAGE_COLUMNS = (LANGUAGE_PERSISTENT, LANGUAGE_EVENTS)
-PERSISTENT_ROW_FIELDS = ("role", "content", "style", "timestamp", "camera", "tool_calls")
-EVENT_ROW_FIELDS = ("role", "content", "style", "camera", "tool_calls")
-
-CORE_STYLES = {
-    "subtask",
-    "plan",
-    "memory",
-    "motion",
-    "interjection",
-    "vqa",
-    "trace",
-    "task_aug",
-}
-EXTENDED_STYLES = set()
-STYLE_REGISTRY = CORE_STYLES | EXTENDED_STYLES
-
-PERSISTENT_STYLES = {"subtask", "plan", "memory", "motion", "task_aug"}
-EVENT_ONLY_STYLES = {"interjection", "vqa", "trace"}
-
-# Styles whose ``content`` is grounded in a specific camera view. Rows of these
-# styles MUST carry a non-null ``camera`` referencing an ``observation.images.*``
-# feature key. Rows of every other style MUST have ``camera=None``. ``motion``
-# is intentionally NOT in this set: motion primitives are described in
-# robot-frame (joint / Cartesian) terms, not pixel space, so they are
-# camera-agnostic. ``trace`` is the pixel-trajectory event style and IS
-# view-dependent. The ``camera`` field nevertheless lives on
-# ``PERSISTENT_ROW_FIELDS`` too so the schema, validator, and resolver
-# behave symmetrically across the two columns; persistent rows simply
-# always have ``camera=None`` in practice today.
-VIEW_DEPENDENT_STYLES = {"vqa", "trace"}
-
-LanguageColumn = Literal["language_persistent", "language_events"]
-
-
-def _json_arrow_type() -> pa.DataType:
-    """Return the Arrow JSON type, falling back to ``string`` on older pyarrow."""
-    return pa.json_() if hasattr(pa, "json_") else pa.string()
-
-
-def _json_feature() -> object:
-    """Return the HF ``datasets`` JSON feature, falling back to a string value."""
-    return datasets.Json() if hasattr(datasets, "Json") else datasets.Value("string")
-
-
-def language_persistent_row_arrow_type() -> pa.StructType:
-    """Return the Arrow struct type for a single persistent language row.
-
-    Persistent rows carry their own ``timestamp`` because they represent a state
-    that became active at a specific moment and remains active until superseded.
-    """
-    return pa.struct(
-        [
-            pa.field("role", pa.string(), nullable=False),
-            pa.field("content", pa.string(), nullable=True),
-            pa.field("style", pa.string(), nullable=True),
-            pa.field("timestamp", pa.float64(), nullable=False),
-            pa.field("camera", pa.string(), nullable=True),
-            pa.field("tool_calls", pa.list_(_json_arrow_type()), nullable=True),
-        ]
-    )
-
-
-def language_event_row_arrow_type() -> pa.StructType:
-    """Return the Arrow struct type for a single event language row.
-
-    Event rows have no ``timestamp`` field: each event is stored on the dataset
-    row whose frame timestamp is the event's firing time.
-    """
-    return pa.struct(
-        [
-            pa.field("role", pa.string(), nullable=False),
-            pa.field("content", pa.string(), nullable=True),
-            pa.field("style", pa.string(), nullable=True),
-            pa.field("camera", pa.string(), nullable=True),
-            pa.field("tool_calls", pa.list_(_json_arrow_type()), nullable=True),
-        ]
-    )
-
-
-def language_persistent_arrow_type() -> pa.ListType:
-    """Return the Arrow list type for the ``language_persistent`` column."""
-    return pa.list_(language_persistent_row_arrow_type())
-
-
-def language_events_arrow_type() -> pa.ListType:
-    """Return the Arrow list type for the ``language_events`` column."""
-    return pa.list_(language_event_row_arrow_type())
-
-
-def language_persistent_row_feature() -> dict[str, object]:
-    """Return the HF ``datasets`` feature mapping for a persistent language row."""
-    return {
-        "role": datasets.Value("string"),
-        "content": datasets.Value("string"),
-        "style": datasets.Value("string"),
-        "timestamp": datasets.Value("float64"),
-        "camera": datasets.Value("string"),
-        "tool_calls": datasets.List(_json_feature()),
-    }
-
-
-def language_event_row_feature() -> dict[str, object]:
-    """Return the HF ``datasets`` feature mapping for an event language row."""
-    return {
-        "role": datasets.Value("string"),
-        "content": datasets.Value("string"),
-        "style": datasets.Value("string"),
-        "camera": datasets.Value("string"),
-        "tool_calls": datasets.List(_json_feature()),
-    }
-
-
-def language_persistent_column_feature() -> datasets.List:
-    """Return the HF ``datasets`` feature for the ``language_persistent`` column."""
-    return datasets.List(language_persistent_row_feature())
-
-
-def language_events_column_feature() -> datasets.List:
-    """Return the HF ``datasets`` feature for the ``language_events`` column."""
-    return datasets.List(language_event_row_feature())
-
-
-def language_feature_info() -> dict[str, dict]:
-    """Return the ``info["features"]`` entries for both language columns."""
-    return {
-        LANGUAGE_PERSISTENT: {"dtype": "language", "shape": (1,), "names": None},
-        LANGUAGE_EVENTS: {"dtype": "language", "shape": (1,), "names": None},
-    }
-
-
-def is_language_column(key: str) -> bool:
-    """Return ``True`` if ``key`` is one of the dataset's language column names."""
-    return key in LANGUAGE_COLUMNS
-
-
-def is_view_dependent_style(style: str | None) -> bool:
-    """Return ``True`` if rows of ``style`` must be tagged with a ``camera`` key."""
-    return style in VIEW_DEPENDENT_STYLES
-
-
-def validate_camera_field(style: str | None, camera: str | None) -> None:
-    """Enforce the ``camera`` invariant: required iff ``style`` is view-dependent.
-
-    Raises ``ValueError`` if a view-dependent style is missing ``camera`` or if
-    a non-view-dependent style carries one. Pipeline writers and the validator
-    should call this on every emitted row.
-    """
-    if is_view_dependent_style(style):
-        if not camera:
-            raise ValueError(
-                f"Rows of view-dependent style {style!r} require a non-empty 'camera' "
-                f"field referencing an 'observation.images.*' feature key."
-            )
-    elif camera is not None:
-        raise ValueError(
-            f"Rows of style {style!r} must have camera=None; got camera={camera!r}."
-        )
-
-
-# --- Tool registry --------------------------------------------------------
-# Tools declared on a dataset live in ``meta/info.json["tools"]`` as a list
-# of OpenAI-style function schemas. The runtime / training stack reads them
-# through :class:`LeRobotDatasetMetadata.tools` (with these constants as
-# fallback when the dataset doesn't declare any). Implementations live
-# under :mod:`lerobot.tools` (one file per tool); see
-# ``docs/source/tools.mdx`` for the authoring guide.
-
-SAY_TOOL_SCHEMA: dict = {
-    "type": "function",
-    "function": {
-        "name": "say",
-        "description": "Speak a short utterance to the user via the TTS executor.",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "text": {
-                    "type": "string",
-                    "description": "The verbatim text to speak.",
-                }
-            },
-            "required": ["text"],
-        },
-    },
-}
-"""Canonical schema for the ``say`` tool emitted by the steerable
-annotation pipeline (PR 2 Module 2). Single source of truth — PR 2's
-writer, PR 3's runtime tool registry, and the dataset visualizer all
-import this constant rather than duplicating the dict."""
-
-DEFAULT_TOOLS: list[dict] = [SAY_TOOL_SCHEMA]
-"""Fallback tools list. Returned by ``LeRobotDatasetMetadata.tools``
-when ``meta/info.json["tools"]`` is unset, so unannotated datasets and
-chat-template consumers (``apply_chat_template(messages, tools=...)``)
-keep working out of the box."""
-
-
-def column_for_style(style: str | None) -> LanguageColumn:
-    """Map a language style to the column where rows of that style are stored.
-
-    Styles in :data:`PERSISTENT_STYLES` route to :data:`LANGUAGE_PERSISTENT`.
-    Styles in :data:`EVENT_ONLY_STYLES` and the implicit ``None`` style route
-    to :data:`LANGUAGE_EVENTS`.
-    """
-    if style is None:
-        return LANGUAGE_EVENTS
-    if style in PERSISTENT_STYLES:
-        return LANGUAGE_PERSISTENT
-    if style in EVENT_ONLY_STYLES:
-        return LANGUAGE_EVENTS
-    raise ValueError(f"Unknown language style: {style!r}")

src/lerobot/datasets/language_render.py

@@ -1,593 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import copy
-import hashlib
-import re
-from collections.abc import Sequence
-from typing import Any
-
-from lerobot.configs.recipe import DEFAULT_BINDINGS, TrainingRecipe
-
-from .language import (
-    EVENT_ONLY_STYLES,
-    LANGUAGE_PERSISTENT,
-    PERSISTENT_STYLES,
-    column_for_style,
-)
-
-LanguageRow = dict[str, Any]
-RenderedMessages = dict[str, list[Any]]
-
-_RESOLVER_RE = re.compile(r"^(?P<name>[A-Za-z_][A-Za-z0-9_]*)\((?P<args>.*)\)$")
-_PLACEHOLDER_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")
-
-
-def active_at(
-    t: float,
-    *,
-    persistent: Sequence[LanguageRow],
-    events: Sequence[LanguageRow] | None = None,
-    style: str | None = None,
-    role: str | None = None,
-    tool_name: str | None = None,
-    camera: str | None = None,
-) -> LanguageRow | None:
-    """Return the persistent row of ``style`` that is active at time ``t``.
-
-    A persistent row is "active" at ``t`` when its own ``timestamp`` is the
-    most recent one ``<= t`` for the given ``style``/``role``/``tool_name``/
-    ``camera`` selector. ``events`` is accepted for resolver-signature
-    uniformity but is not consulted: only persistent styles are valid here.
-    """
-    _validate_persistent_resolver("active_at", style)
-    matches = _matching_rows(
-        persistent, style=style, role=role, tool_name=tool_name, camera=camera
-    )
-    matches = [row for row in matches if _timestamp(row) <= t]
-    return _select_latest(
-        matches, style=style, role=role, tool_name=tool_name, camera=camera
-    )
-
-
-def emitted_at(
-    t: float,
-    *,
-    persistent: Sequence[LanguageRow],
-    events: Sequence[LanguageRow],
-    style: str | None = None,
-    role: str | None = None,
-    tool_name: str | None = None,
-    camera: str | None = None,
-) -> LanguageRow | None:
-    """Return the row of ``style`` emitted at exactly time ``t``.
-
-    For persistent styles, this matches persistent rows whose own ``timestamp``
-    equals ``t``. For event styles, the ``events`` list is assumed to come from
-    the dataset row at frame ``t`` (event rows carry no timestamp of their own),
-    so all matching event rows are considered emitted at ``t``. ``camera``
-    filters by the row's ``camera`` field — required to disambiguate when
-    multiple view-dependent rows share ``(t, role)`` across cameras.
-    """
-    column = column_for_style(style)
-    if column == LANGUAGE_PERSISTENT:
-        matches = [
-            row
-            for row in _matching_rows(
-                persistent, style=style, role=role, tool_name=tool_name, camera=camera
-            )
-            if _timestamp(row) == t
-        ]
-        return _select_one(
-            matches,
-            style=style,
-            role=role,
-            tool_name=tool_name,
-            camera=camera,
-            sort_key=_persistent_sort_key,
-        )
-    matches = _matching_rows(
-        events, style=style, role=role, tool_name=tool_name, camera=camera
-    )
-    return _select_one(
-        matches,
-        style=style,
-        role=role,
-        tool_name=tool_name,
-        camera=camera,
-        sort_key=_event_sort_key,
-    )
-
-
-def nth_prev(
-    t: float,
-    *,
-    persistent: Sequence[LanguageRow],
-    events: Sequence[LanguageRow] | None = None,
-    style: str | None = None,
-    offset: int = 1,
-    role: str | None = None,
-    tool_name: str | None = None,
-    camera: str | None = None,
-) -> LanguageRow | None:
-    """Return the persistent row that was active ``offset`` steps before ``t``.
-
-    Walks back through chronologically sorted persistent rows of ``style``
-    (filtered by optional ``role``/``tool_name``/``camera``) and returns the
-    one ``offset`` positions before the row active at ``t``. Only valid for
-    persistent styles.
-    """
-    return _nth_relative(
-        t,
-        persistent=persistent,
-        style=style,
-        offset=-offset,
-        role=role,
-        tool_name=tool_name,
-        camera=camera,
-        resolver_name="nth_prev",
-    )
-
-
-def nth_next(
-    t: float,
-    *,
-    persistent: Sequence[LanguageRow],
-    events: Sequence[LanguageRow] | None = None,
-    style: str | None = None,
-    offset: int = 1,
-    role: str | None = None,
-    tool_name: str | None = None,
-    camera: str | None = None,
-) -> LanguageRow | None:
-    """Return the persistent row that becomes active ``offset`` steps after ``t``.
-
-    Walks forward through chronologically sorted persistent rows of ``style``
-    (filtered by optional ``role``/``tool_name``/``camera``) and returns the
-    one ``offset`` positions after the row active at ``t``. Only valid for
-    persistent styles.
-    """
-    return _nth_relative(
-        t,
-        persistent=persistent,
-        style=style,
-        offset=offset,
-        role=role,
-        tool_name=tool_name,
-        camera=camera,
-        resolver_name="nth_next",
-    )
-
-
-def render_sample(
-    *,
-    recipe: TrainingRecipe,
-    persistent: Sequence[LanguageRow] | None,
-    events: Sequence[LanguageRow] | None,
-    t: float,
-    sample_idx: int,
-    task: str | None = None,
-    dataset_ctx: Any | None = None,
-) -> RenderedMessages | None:
-    """Render the chat-style messages for a single dataset sample.
-
-    Resolves the recipe's bindings against ``persistent`` and ``events`` rows
-    at frame timestamp ``t``, then expands the recipe's message templates.
-    Returns ``None`` if the resolved sample contains no target message.
-    """
-    persistent_rows = _normalize_rows(persistent or [])
-    event_rows = _normalize_rows(events or [])
-    selected_recipe = _select_recipe(recipe, sample_idx)
-    bindings = _resolve_bindings(
-        selected_recipe,
-        persistent=persistent_rows,
-        events=event_rows,
-        t=t,
-        sample_idx=sample_idx,
-        task=task,
-        dataset_ctx=dataset_ctx,
-    )
-    return _render_message_recipe(selected_recipe, bindings)
-
-
-def _select_recipe(recipe: TrainingRecipe, sample_idx: int) -> TrainingRecipe:
-    """Pick a deterministic blend component for ``sample_idx`` (or return ``recipe``)."""
-    if recipe.blend is None:
-        return recipe
-
-    total_weight = sum(component.weight or 0.0 for component in recipe.blend.values())
-    if total_weight <= 0:
-        raise ValueError("Blend weights must sum to a positive value.")
-
-    digest = hashlib.blake2b(str(sample_idx).encode(), digest_size=8).digest()
-    draw = int.from_bytes(digest, "big") / 2**64 * total_weight
-    cumulative = 0.0
-    last_component: TrainingRecipe | None = None
-    for component in recipe.blend.values():
-        last_component = component
-        cumulative += component.weight or 0.0
-        if draw < cumulative:
-            return component
-    assert last_component is not None
-    return last_component
-
-
-def _resolve_bindings(
-    recipe: TrainingRecipe,
-    *,
-    persistent: Sequence[LanguageRow],
-    events: Sequence[LanguageRow],
-    t: float,
-    sample_idx: int,
-    task: str | None,
-    dataset_ctx: Any | None,
-) -> dict[str, LanguageRow | str | None]:
-    """Resolve every binding in ``recipe`` (plus ``task``) at time ``t``."""
-    bindings: dict[str, LanguageRow | str | None] = {
-        "task": _resolve_task(
-            task, dataset_ctx, persistent=persistent, sample_idx=sample_idx
-        ),
-    }
-    specs = {**DEFAULT_BINDINGS, **(recipe.bindings or {})}
-    for name, spec in specs.items():
-        bindings[name] = _resolve_spec(spec, persistent=persistent, events=events, t=t)
-    return bindings
-
-
-def _resolve_task(
-    task: str | None,
-    dataset_ctx: Any | None,
-    *,
-    persistent: Sequence[LanguageRow] = (),
-    sample_idx: int = 0,
-) -> str | None:
-    """Return the task string for ``sample_idx``.
-
-    Resolution order:
-
-    1. Explicit ``task`` override (caller-supplied) wins.
-    2. If ``persistent`` contains rows of style ``task_aug`` (role=user),
-       deterministically pick one by ``sample_idx`` so each frame of an
-       episode rotates through the available rephrasings across an epoch.
-       This realizes Xiao 2022 / CAST-style task-prompt diversity without
-       changing ``meta/tasks.parquet`` and without forcing recipes to opt
-       in: ``${task}`` automatically picks a rephrasing when one exists,
-       and falls back to the canonical task otherwise. Recipes that want
-       the literal canonical task can override the binding.
-    3. Otherwise read the canonical task from ``dataset_ctx`` (which is
-       backed by ``meta/tasks.parquet``).
-    """
-    if task is not None:
-        return task
-
-    aug_rows = [
-        r
-        for r in persistent
-        if r.get("style") == "task_aug" and r.get("role") == "user"
-    ]
-    if aug_rows:
-        # Deterministic, blake2b-based pick keyed on sample_idx so the
-        # rotation is reproducible across runs (Python's built-in ``hash``
-        # is process-randomized).
-        digest = hashlib.blake2b(
-            f"task_aug:{sample_idx}".encode(), digest_size=8
-        ).digest()
-        idx = int.from_bytes(digest, "big") % len(aug_rows)
-        chosen = aug_rows[idx].get("content")
-        if chosen:
-            return str(chosen)
-
-    if dataset_ctx is None:
-        return None
-    if isinstance(dataset_ctx, dict):
-        return dataset_ctx.get("task")
-    return getattr(dataset_ctx, "task", None)
-
-
-def _resolve_spec(
-    spec: str,
-    *,
-    persistent: Sequence[LanguageRow],
-    events: Sequence[LanguageRow],
-    t: float,
-) -> LanguageRow | None:
-    """Parse a single binding's resolver expression and dispatch to its function."""
-    match = _RESOLVER_RE.match(spec.strip())
-    if match is None:
-        raise ValueError(f"Invalid resolver expression: {spec!r}")
-    name = match.group("name")
-    kwargs = _parse_resolver_args(match.group("args"))
-    kwargs.pop("t_arg", None)
-
-    resolvers = {
-        "active_at": active_at,
-        "emitted_at": emitted_at,
-        "nth_prev": nth_prev,
-        "nth_next": nth_next,
-    }
-    if name not in resolvers:
-        raise ValueError(f"Unknown language resolver: {name!r}")
-    return resolvers[name](t, persistent=persistent, events=events, **kwargs)
-
-
-def _parse_resolver_args(args: str) -> dict[str, Any]:
-    """Parse a comma-separated resolver argument list into a kwargs dict."""
-    kwargs: dict[str, Any] = {}
-    if not args.strip():
-        return kwargs
-
-    parts = [part.strip() for part in args.split(",") if part.strip()]
-    for part in parts:
-        if part == "t":
-            kwargs["t_arg"] = True
-            continue
-        if "=" not in part:
-            raise ValueError(f"Invalid resolver argument: {part!r}")
-        key, value = (item.strip() for item in part.split("=", 1))
-        if key == "offset":
-            kwargs[key] = int(value)
-        else:
-            kwargs[key] = value.strip("\"'")
-    return kwargs
-
-
-def _render_message_recipe(
-    recipe: TrainingRecipe,
-    bindings: dict[str, LanguageRow | str | None],
-) -> RenderedMessages | None:
-    """Expand ``recipe.messages`` into rendered chat messages using ``bindings``."""
-    assert recipe.messages is not None
-    messages: list[dict[str, Any]] = []
-    streams: list[str | None] = []
-    target_indices: list[int] = []
-
-    for turn in recipe.messages:
-        if turn.if_present is not None and bindings.get(turn.if_present) is None:
-            continue
-
-        message = {"role": turn.role}
-        if turn.content is not None:
-            message["content"] = _render_content(turn.content, bindings)
-
-        if turn.tool_calls_from is not None:
-            row = bindings.get(turn.tool_calls_from)
-            tool_calls = row.get("tool_calls") if isinstance(row, dict) else None
-            if tool_calls:
-                message["tool_calls"] = copy.deepcopy(tool_calls)
-
-        message_idx = len(messages)
-        messages.append(message)
-        streams.append(turn.stream)
-        if turn.target:
-            target_indices.append(message_idx)
-
-    if not target_indices:
-        return None
-
-    rendered = {
-        "messages": messages,
-        "message_streams": streams,
-        "target_message_indices": target_indices,
-    }
-    _validate_rendered(rendered)
-    return rendered
-
-
-def _render_content(
-    content: str | list[dict[str, Any]],
-    bindings: dict[str, LanguageRow | str | None],
-) -> str | list[dict[str, Any]]:
-    """Substitute bindings into a string or each string field of multimodal blocks."""
-    if isinstance(content, str):
-        return _substitute(content, bindings)
-
-    rendered_blocks = []
-    for block in content:
-        rendered_block = copy.deepcopy(block)
-        for key, value in rendered_block.items():
-            if isinstance(value, str):
-                rendered_block[key] = _substitute(value, bindings)
-        rendered_blocks.append(rendered_block)
-    return rendered_blocks
-
-
-def _substitute(template: str, bindings: dict[str, LanguageRow | str | None]) -> str:
-    """Replace ``${name}`` placeholders in ``template`` with their bound values."""
-
-    def replace(match: re.Match[str]) -> str:
-        """Resolve a single ``${name}`` match to its bound string value."""
-        name = match.group(1)
-        if name not in bindings:
-            raise ValueError(f"Unknown template binding: {name!r}")
-        value = bindings[name]
-        if value is None:
-            return ""
-        if isinstance(value, dict):
-            content = value.get("content")
-            return "" if content is None else str(content)
-        return str(value)
-
-    return _PLACEHOLDER_RE.sub(replace, template)
-
-
-def _validate_rendered(rendered: RenderedMessages) -> None:
-    """Sanity-check the rendered output for stream/target alignment."""
-    messages = rendered["messages"]
-    streams = rendered["message_streams"]
-    target_indices = rendered["target_message_indices"]
-
-    if len(streams) != len(messages):
-        raise ValueError("message_streams must be aligned with messages.")
-    if not target_indices:
-        raise ValueError("Rendered samples must contain at least one target message.")
-    for idx in target_indices:
-        if idx < 0 or idx >= len(messages):
-            raise ValueError(f"Target message index {idx} is out of bounds.")
-    for idx, stream in enumerate(streams):
-        if stream is None:
-            raise ValueError(f"Rendered message {idx} has no stream.")
-
-
-def _nth_relative(
-    t: float,
-    *,
-    persistent: Sequence[LanguageRow],
-    style: str | None,
-    offset: int,
-    role: str | None,
-    tool_name: str | None,
-    camera: str | None,
-    resolver_name: str,
-) -> LanguageRow | None:
-    """Shared body for ``nth_prev`` / ``nth_next`` with signed ``offset``."""
-    _validate_persistent_resolver(resolver_name, style)
-    if abs(offset) < 1:
-        raise ValueError(f"{resolver_name} offset must be non-zero.")
-
-    rows = sorted(
-        _matching_rows(persistent, style=style, role=role, tool_name=tool_name, camera=camera),
-        key=_persistent_sort_key,
-    )
-    if not rows:
-        return None
-
-    anchor_idx = None
-    for idx, row in enumerate(rows):
-        if _timestamp(row) <= t:
-            anchor_idx = idx
-        else:
-            break
-
-    target_idx = (offset - 1 if offset > 0 else None) if anchor_idx is None else anchor_idx + offset
-
-    if target_idx is None or target_idx < 0 or target_idx >= len(rows):
-        return None
-    return rows[target_idx]
-
-
-def _validate_persistent_resolver(resolver_name: str, style: str | None) -> None:
-    """Reject calls with missing or event-only ``style`` for persistent resolvers."""
-    if style is None:
-        raise ValueError(f"{resolver_name} requires a persistent style.")
-    if style in EVENT_ONLY_STYLES:
-        raise ValueError(f"{resolver_name} cannot be used with event-only style {style!r}.")
-    if style not in PERSISTENT_STYLES:
-        column_for_style(style)
-
-
-def _matching_rows(
-    rows: Sequence[LanguageRow],
-    *,
-    style: str | None,
-    role: str | None,
-    tool_name: str | None,
-    camera: str | None,
-) -> list[LanguageRow]:
-    """Return ``rows`` filtered by optional ``style``/``role``/``tool_name``/``camera`` selectors."""
-    return [
-        row
-        for row in rows
-        if (style is None or row.get("style") == style)
-        and (role is None or row.get("role") == role)
-        and (tool_name is None or _row_has_tool_name(row, tool_name))
-        and (camera is None or row.get("camera") == camera)
-    ]
-
-
-def _select_latest(
-    rows: Sequence[LanguageRow],
-    *,
-    style: str | None,
-    role: str | None,
-    tool_name: str | None,
-    camera: str | None,
-) -> LanguageRow | None:
-    """Return the row tied for the latest ``timestamp`` (disambiguated by selectors)."""
-    if not rows:
-        return None
-    rows = sorted(rows, key=_persistent_sort_key)
-    latest_ts = _timestamp(rows[-1])
-    return _select_one(
-        [row for row in rows if _timestamp(row) == latest_ts],
-        style=style,
-        role=role,
-        tool_name=tool_name,
-        camera=camera,
-        sort_key=_persistent_sort_key,
-    )
-
-
-def _select_one(
-    rows: Sequence[LanguageRow],
-    *,
-    style: str | None,
-    role: str | None,
-    tool_name: str | None,
-    camera: str | None,
-    sort_key: Any,
-) -> LanguageRow | None:
-    """Return the single matching row, or raise if the selectors are ambiguous."""
-    if not rows:
-        return None
-    if len(rows) > 1 and role is None and tool_name is None and camera is None:
-        raise ValueError(
-            f"Ambiguous resolver for style={style!r}; add role=..., tool_name=..., "
-            f"or camera=... to disambiguate."
-        )
-    return sorted(rows, key=sort_key)[0]
-
-
-def _persistent_sort_key(row: LanguageRow) -> tuple[float, str, str]:
-    """Sort key for persistent rows: ``(timestamp, style, role)``."""
-    return (_timestamp(row), row.get("style") or "", row.get("role") or "")
-
-
-def _event_sort_key(row: LanguageRow) -> tuple[str, str]:
-    """Sort key for event rows: ``(style, role)`` (timestamp is implicit in the frame)."""
-    return (row.get("style") or "", row.get("role") or "")
-
-
-def _timestamp(row: LanguageRow) -> float:
-    """Extract a row's ``timestamp`` as a Python float (unwrapping numpy scalars)."""
-    value = row["timestamp"]
-    return float(value.item() if hasattr(value, "item") else value)
-
-
-def _row_has_tool_name(row: LanguageRow, tool_name: str) -> bool:
-    """Return ``True`` if any of the row's tool calls invokes ``tool_name``."""
-    for tool_call in row.get("tool_calls") or []:
-        if isinstance(tool_call, str):
-            continue
-        function = tool_call.get("function") if isinstance(tool_call, dict) else None
-        if isinstance(function, dict) and function.get("name") == tool_name:
-            return True
-    return False
-
-
-def _normalize_rows(rows: Sequence[Any]) -> list[LanguageRow]:
-    """Convert pyarrow scalars / mappings into a fresh list of plain dict rows."""
-    normalized = []
-    for row in rows:
-        if row is None:
-            continue
-        if hasattr(row, "as_py"):
-            row = row.as_py()
-        if not isinstance(row, dict):
-            raise TypeError(f"Language rows must be dictionaries, got {type(row).__name__}.")
-        normalized.append(dict(row))
-    return normalized

src/lerobot/datasets/utils.py

@@ -71,8 +71,8 @@ class ForwardCompatibilityError(CompatibilityError):
 
 
 DEFAULT_CHUNK_SIZE = 1000  # Max number of files per chunk
-DEFAULT_DATA_FILE_SIZE_IN_MB = 100  # Max size per file
-DEFAULT_VIDEO_FILE_SIZE_IN_MB = 200  # Max size per file
+DEFAULT_DATA_FILE_SIZE_IN_MB = 50  # Max size per file
+DEFAULT_VIDEO_FILE_SIZE_IN_MB = 100  # Max size per file
 
 INFO_PATH = "meta/info.json"
 STATS_PATH = "meta/stats.json"
@@ -83,6 +83,7 @@ VIDEO_DIR = "videos"
 
 CHUNK_FILE_PATTERN = "chunk-{chunk_index:03d}/file-{file_index:03d}"
 DEFAULT_TASKS_PATH = "meta/tasks.parquet"
+DEFAULT_SUBTASKS_PATH = "meta/subtasks.parquet"
 DEFAULT_EPISODES_PATH = EPISODES_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet"
 DEFAULT_DATA_PATH = DATA_DIR + "/" + CHUNK_FILE_PATTERN + ".parquet"
 DEFAULT_VIDEO_PATH = VIDEO_DIR + "/{video_key}/" + CHUNK_FILE_PATTERN + ".mp4"

src/lerobot/envs/configs.py

@@ -331,7 +331,6 @@ class LiberoEnv(EnvConfig):
     camera_name_mapping: dict[str, str] | None = None
     observation_height: int = 360
     observation_width: int = 360
-    is_libero_plus: bool = False
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
             ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
@@ -433,7 +432,6 @@ class LiberoEnv(EnvConfig):
             control_mode=self.control_mode,
             episode_length=self.episode_length,
             camera_name_mapping=self.camera_name_mapping,
-            is_libero_plus=self.is_libero_plus,
         )
 
     def get_env_processors(self):
@@ -498,146 +496,6 @@ class MetaworldEnv(EnvConfig):
         )
 
 
-@EnvConfig.register_subclass("robocasa")
-@dataclass
-class RoboCasaEnv(EnvConfig):
-    task: str = "CloseFridge"
-    fps: int = 20
-    episode_length: int = 1000
-    obs_type: str = "pixels_agent_pos"
-    render_mode: str = "rgb_array"
-    camera_name: str = "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right"
-    observation_height: int = 256
-    observation_width: int = 256
-    visualization_height: int = 512
-    visualization_width: int = 512
-    split: str | None = None
-    # Object-mesh registries to sample from. Upstream default is
-    # ("objaverse", "lightwheel"), but objaverse is ~30GB and the CI image
-    # only ships the lightwheel pack. Override to include objaverse once
-    # you've run `python -m robocasa.scripts.download_kitchen_assets
-    # --type objaverse` locally.
-    obj_registries: list[str] = field(default_factory=lambda: ["lightwheel"])
-    features: dict[str, PolicyFeature] = field(
-        default_factory=lambda: {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(12,))}
-    )
-    features_map: dict[str, str] = field(default_factory=lambda: {ACTION: ACTION, "agent_pos": OBS_STATE})
-
-    def __post_init__(self):
-        if self.obs_type not in ("pixels", "pixels_agent_pos"):
-            raise ValueError(f"Unsupported obs_type: {self.obs_type}")
-
-        # Preserve raw RoboCasa camera names end-to-end (e.g.
-        # `observation.images.robot0_agentview_left`). This matches the
-        # naming convention used by the RoboCasa datasets on the Hub, so
-        # trained policies don't need a `--rename_map` at eval time.
-        cams = [c.strip() for c in self.camera_name.split(",") if c.strip()]
-        for cam in cams:
-            self.features[f"pixels/{cam}"] = PolicyFeature(
-                type=FeatureType.VISUAL,
-                shape=(self.observation_height, self.observation_width, 3),
-            )
-            self.features_map[f"pixels/{cam}"] = f"{OBS_IMAGES}.{cam}"
-
-        if self.obs_type == "pixels_agent_pos":
-            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(16,))
-
-    @property
-    def gym_kwargs(self) -> dict:
-        kwargs: dict[str, Any] = {
-            "obs_type": self.obs_type,
-            "render_mode": self.render_mode,
-            "observation_height": self.observation_height,
-            "observation_width": self.observation_width,
-            "visualization_height": self.visualization_height,
-            "visualization_width": self.visualization_width,
-        }
-        if self.split is not None:
-            kwargs["split"] = self.split
-        return kwargs
-
-    def create_envs(self, n_envs: int, use_async_envs: bool = False):
-        from .robocasa import create_robocasa_envs
-
-        if self.task is None:
-            raise ValueError("RoboCasaEnv requires a task to be specified")
-        env_cls = _make_vec_env_cls(use_async_envs, n_envs)
-        return create_robocasa_envs(
-            task=self.task,
-            n_envs=n_envs,
-            camera_name=self.camera_name,
-            gym_kwargs=self.gym_kwargs,
-            env_cls=env_cls,
-            episode_length=self.episode_length,
-            obj_registries=tuple(self.obj_registries),
-        )
-
-
-@EnvConfig.register_subclass("vlabench")
-@dataclass
-class VLABenchEnv(EnvConfig):
-    task: str = "select_fruit"
-    fps: int = 10
-    episode_length: int = 500
-    obs_type: str = "pixels_agent_pos"
-    render_mode: str = "rgb_array"
-    render_resolution: tuple[int, int] = (480, 480)
-    robot: str = "franka"
-    action_mode: str = "eef"
-    features: dict[str, PolicyFeature] = field(
-        default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
-        }
-    )
-    features_map: dict[str, str] = field(
-        default_factory=lambda: {
-            ACTION: ACTION,
-            "agent_pos": OBS_STATE,
-            "pixels/image": f"{OBS_IMAGES}.image",
-            "pixels/second_image": f"{OBS_IMAGES}.second_image",
-            "pixels/wrist_image": f"{OBS_IMAGES}.wrist_image",
-        }
-    )
-
-    def __post_init__(self):
-        h, w = self.render_resolution
-        if self.obs_type == "pixels":
-            self.features["pixels/image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3))
-            self.features["pixels/second_image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3))
-            self.features["pixels/wrist_image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3))
-        elif self.obs_type == "pixels_agent_pos":
-            self.features["pixels/image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3))
-            self.features["pixels/second_image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3))
-            self.features["pixels/wrist_image"] = PolicyFeature(type=FeatureType.VISUAL, shape=(h, w, 3))
-            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(7,))
-        else:
-            raise ValueError(f"Unsupported obs_type: {self.obs_type}")
-
-    @property
-    def gym_kwargs(self) -> dict:
-        return {
-            "obs_type": self.obs_type,
-            "render_mode": self.render_mode,
-            "render_resolution": self.render_resolution,
-            "robot": self.robot,
-            "max_episode_steps": self.episode_length,
-            "action_mode": self.action_mode,
-        }
-
-    def create_envs(self, n_envs: int, use_async_envs: bool = False):
-        from .vlabench import create_vlabench_envs
-
-        if self.task is None:
-            raise ValueError("VLABenchEnv requires a task to be specified")
-        env_cls = _make_vec_env_cls(use_async_envs, n_envs)
-        return create_vlabench_envs(
-            task=self.task,
-            n_envs=n_envs,
-            gym_kwargs=self.gym_kwargs,
-            env_cls=env_cls,
-        )
-
-
 @EnvConfig.register_subclass("isaaclab_arena")
 @dataclass
 class IsaaclabArenaEnv(HubEnvConfig):
@@ -716,171 +574,3 @@ class IsaaclabArenaEnv(HubEnvConfig):
             ),
             PolicyProcessorPipeline(steps=[]),
         )
-
-
-@EnvConfig.register_subclass("libero_plus")
-@dataclass
-class LiberoPlusEnv(LiberoEnv):
-    """Config for LIBERO-plus robustness benchmark evaluation.
-
-    LIBERO-plus extends LIBERO with 7 perturbation dimensions (camera viewpoints,
-    object layouts, robot initial states, language instructions, lighting, background
-    textures, sensor noise) producing ~10k task variants.
-
-    The gym interface is identical to LIBERO so this class reuses ``LiberoEnv``
-    entirely — only the registered name and default task suite differ.
-
-    Install: see docker/Dockerfile.benchmark.libero_plus — LIBERO-plus ships
-    as a namespace package from a git fork and must be cloned + PYTHONPATH'd
-    rather than installed as a pyproject extra.
-
-    See Also:
-        https://github.com/sylvestf/LIBERO-plus
-    """
-
-    task: str = "libero_spatial"
-    is_libero_plus: bool = True
-
-
-@EnvConfig.register_subclass("robotwin")
-@dataclass
-class RoboTwinEnvConfig(EnvConfig):
-    """Configuration for RoboTwin 2.0 benchmark environments.
-
-    RoboTwin 2.0 is a dual-arm manipulation benchmark with 50 tasks built on the
-    SAPIEN simulator. The robot is an Aloha-AgileX bimanual platform with 14 DOF
-    (7 per arm). All three cameras are enabled by default.
-
-    See: https://robotwin-platform.github.io
-    Dataset: https://huggingface.co/datasets/lerobot/robotwin_unified
-    """
-
-    task: str = "beat_block_hammer"  # single task or comma-separated list
-    fps: int = 25
-    episode_length: int = 300
-    obs_type: str = "pixels_agent_pos"
-    render_mode: str = "rgb_array"
-    # Available cameras from RoboTwin's aloha-agilex embodiment: head_camera
-    # (torso-mounted) + left_camera / right_camera (wrists).
-    camera_names: str = "head_camera,left_camera,right_camera"
-    # Match the D435 dims in task_config/demo_clean.yml (_camera_config.yml).
-    # Gym's vector-env concatenate pre-allocates buffers of this shape, so it
-    # must equal what SAPIEN actually renders.
-    observation_height: int = 240
-    observation_width: int = 320
-    features: dict[str, PolicyFeature] = field(
-        default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(14,)),
-        }
-    )
-    features_map: dict[str, str] = field(
-        default_factory=lambda: {
-            ACTION: ACTION,
-            "pixels/head_camera": f"{OBS_IMAGES}.head_camera",
-            "pixels/left_camera": f"{OBS_IMAGES}.left_camera",
-            "pixels/right_camera": f"{OBS_IMAGES}.right_camera",
-            "agent_pos": OBS_STATE,
-        }
-    )
-
-    def __post_init__(self):
-        cam_list = [c.strip() for c in self.camera_names.split(",") if c.strip()]
-        for cam in cam_list:
-            self.features[f"pixels/{cam}"] = PolicyFeature(
-                type=FeatureType.VISUAL,
-                shape=(self.observation_height, self.observation_width, 3),
-            )
-            # Keep features_map entry if already set (default_factory); add if missing.
-            key = f"pixels/{cam}"
-            if key not in self.features_map:
-                self.features_map[key] = f"{OBS_IMAGES}.{cam}"
-
-        if self.obs_type == "pixels_agent_pos":
-            self.features["agent_pos"] = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(14,),  # 14 DOF: 7 per arm
-            )
-        elif self.obs_type != "pixels":
-            raise ValueError(
-                f"Unsupported obs_type '{self.obs_type}'. "
-                "RoboTwinEnvConfig supports 'pixels' and 'pixels_agent_pos'."
-            )
-
-    @property
-    def gym_kwargs(self) -> dict:
-        return {}
-
-    def create_envs(self, n_envs: int, use_async_envs: bool = True):
-        from lerobot.envs.robotwin import create_robotwin_envs
-
-        if not self.task:
-            raise ValueError("RoboTwinEnvConfig requires `task` to be specified.")
-
-        env_cls = _make_vec_env_cls(use_async_envs, n_envs)
-        cam_list = [c.strip() for c in self.camera_names.split(",") if c.strip()]
-        return create_robotwin_envs(
-            task=self.task,
-            n_envs=n_envs,
-            env_cls=env_cls,
-            camera_names=cam_list,
-            observation_height=self.observation_height,
-            observation_width=self.observation_width,
-            episode_length=self.episode_length,
-        )
-
-
-@EnvConfig.register_subclass("robomme")
-@dataclass
-class RoboMMEEnv(EnvConfig):
-    """RoboMME memory-augmented manipulation benchmark (ManiSkill/SAPIEN).
-
-    16 tasks across 4 suites: Counting, Permanence, Reference, Imitation.
-    Dataset: lerobot/robomme (LeRobot v3.0, 1,600 episodes).
-    Benchmark: https://github.com/RoboMME/robomme_benchmark
-
-    Requires the `robomme` git package installed separately (Linux only);
-    see docker/Dockerfile.benchmark.robomme for the canonical install.
-    """
-
-    task: str = "PickXtimes"
-    fps: int = 10
-    episode_length: int = 300
-    action_space: str = "joint_angle"  # or "ee_pose" (7-D)
-    dataset_split: str = "test"  # "train" | "val" | "test"
-    task_ids: list[int] | None = None
-    features: dict[str, PolicyFeature] = field(default_factory=dict)
-    features_map: dict[str, str] = field(
-        default_factory=lambda: {
-            ACTION: ACTION,
-            "pixels/image": f"{OBS_IMAGES}.image",
-            "pixels/wrist_image": f"{OBS_IMAGES}.wrist_image",
-            "agent_pos": OBS_STATE,
-        }
-    )
-
-    def __post_init__(self):
-        action_dim = 8 if self.action_space == "joint_angle" else 7
-        self.features = {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,)),
-            "pixels/image": PolicyFeature(type=FeatureType.VISUAL, shape=(256, 256, 3)),
-            "pixels/wrist_image": PolicyFeature(type=FeatureType.VISUAL, shape=(256, 256, 3)),
-            "agent_pos": PolicyFeature(type=FeatureType.STATE, shape=(8,)),
-        }
-
-    @property
-    def gym_kwargs(self) -> dict:
-        return {}
-
-    def create_envs(self, n_envs: int, use_async_envs: bool = True):
-        from lerobot.envs.robomme import create_robomme_envs
-
-        env_cls = _make_vec_env_cls(use_async_envs, n_envs)
-        return create_robomme_envs(
-            task=self.task,
-            n_envs=n_envs,
-            action_space_type=self.action_space,
-            dataset=self.dataset_split,
-            episode_length=self.episode_length,
-            task_ids=self.task_ids,
-            env_cls=env_cls,
-        )

src/lerobot/envs/libero.py

@@ -16,7 +16,6 @@
 from __future__ import annotations
 
 import os
-import re
 from collections import defaultdict
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from functools import partial
@@ -32,7 +31,20 @@ from libero.libero.envs import OffScreenRenderEnv
 
 from lerobot.types import RobotObservation
 
-from .utils import _LazyAsyncVectorEnv, parse_camera_names
+from .utils import _LazyAsyncVectorEnv
+
+
+def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
+    """Normalize camera_name into a non-empty list of strings."""
+    if isinstance(camera_name, str):
+        cams = [c.strip() for c in camera_name.split(",") if c.strip()]
+    elif isinstance(camera_name, (list | tuple)):
+        cams = [str(c).strip() for c in camera_name if str(c).strip()]
+    else:
+        raise TypeError(f"camera_name must be str or sequence[str], got {type(camera_name).__name__}")
+    if not cams:
+        raise ValueError("camera_name resolved to an empty list.")
+    return cams
 
 
 def _get_suite(name: str) -> benchmark.Benchmark:
@@ -57,34 +69,14 @@ def _select_task_ids(total_tasks: int, task_ids: Iterable[int] | None) -> list[i
     return ids
 
 
-# LIBERO-plus perturbation variants encode the perturbation in the filename
-# but on disk only the base `.pruned_init` exists — strip the suffix to match
-# LIBERO-plus's own suite.get_task_init_states() (we reimplement it here so we
-# can pass weights_only=False for PyTorch 2.6+ numpy pickles).
-_LIBERO_PERTURBATION_SUFFIX_RE = re.compile(r"_(?:language|view|light)_[^.]*|_(?:table|tb)_\d+")
-
-
-def get_task_init_states(task_suite: Any, i: int, is_libero_plus: bool = False) -> np.ndarray:
-    task = task_suite.tasks[i]
-    filename = Path(task.init_states_file)
-    root = Path(get_libero_path("init_states"))
-
-    if not is_libero_plus:
-        init_states_path = root / task.problem_folder / filename.name
-        return torch.load(init_states_path, weights_only=False)  # nosec B614
-
-    # LIBERO-plus: `_add_` / `_level` variants store extra-object layouts under
-    # libero_newobj/ as a flat array that must be reshaped to (1, -1).
-    if "_add_" in filename.name or "_level" in filename.name:
-        init_states_path = root / "libero_newobj" / task.problem_folder / filename.name
-        init_states = torch.load(init_states_path, weights_only=False)  # nosec B614
-        return init_states.reshape(1, -1)
-
-    # LIBERO-plus perturbation variants encode the perturbation in the filename
-    # but on disk only the base `.pruned_init` exists — strip the suffix to match.
-    stripped = _LIBERO_PERTURBATION_SUFFIX_RE.sub("", filename.stem) + filename.suffix
-    init_states_path = root / task.problem_folder / stripped
-    return torch.load(init_states_path, weights_only=False)  # nosec B614
+def get_task_init_states(task_suite: Any, i: int) -> np.ndarray:
+    init_states_path = (
+        Path(get_libero_path("init_states"))
+        / task_suite.tasks[i].problem_folder
+        / task_suite.tasks[i].init_states_file
+    )
+    init_states = torch.load(init_states_path, weights_only=False)  # nosec B614
+    return init_states
 
 
 def get_libero_dummy_action():
@@ -126,11 +118,9 @@ class LiberoEnv(gym.Env):
         camera_name_mapping: dict[str, str] | None = None,
         num_steps_wait: int = 10,
         control_mode: str = "relative",
-        is_libero_plus: bool = False,
     ):
         super().__init__()
         self.task_id = task_id
-        self.is_libero_plus = is_libero_plus
         self.obs_type = obs_type
         self.render_mode = render_mode
         self.observation_width = observation_width
@@ -138,7 +128,7 @@ class LiberoEnv(gym.Env):
         self.visualization_width = visualization_width
         self.visualization_height = visualization_height
         self.init_states = init_states
-        self.camera_name = parse_camera_names(
+        self.camera_name = _parse_camera_names(
             camera_name
         )  # agentview_image (main) or robot0_eye_in_hand_image (wrist)
 
@@ -157,11 +147,7 @@ class LiberoEnv(gym.Env):
         self.episode_index = episode_index
         self.episode_length = episode_length
         # Load once and keep
-        self._init_states = (
-            get_task_init_states(task_suite, self.task_id, is_libero_plus=self.is_libero_plus)
-            if self.init_states
-            else None
-        )
+        self._init_states = get_task_init_states(task_suite, self.task_id) if self.init_states else None
         self._reset_stride = n_envs  # when performing a reset, append `_reset_stride` to `init_state_id`.
 
         self.init_state_id = self.episode_index  # tie each sub-env to a fixed init state
@@ -394,7 +380,6 @@ def _make_env_fns(
     gym_kwargs: Mapping[str, Any],
     control_mode: str,
     camera_name_mapping: dict[str, str] | None = None,
-    is_libero_plus: bool = False,
 ) -> list[Callable[[], LiberoEnv]]:
     """Build n_envs factory callables for a single (suite, task_id)."""
 
@@ -411,7 +396,6 @@ def _make_env_fns(
             n_envs=n_envs,
             control_mode=control_mode,
             camera_name_mapping=camera_name_mapping,
-            is_libero_plus=is_libero_plus,
             **local_kwargs,
         )
 
@@ -434,7 +418,6 @@ def create_libero_envs(
     control_mode: str = "relative",
     episode_length: int | None = None,
     camera_name_mapping: dict[str, str] | None = None,
-    is_libero_plus: bool = False,
 ) -> dict[str, dict[int, Any]]:
     """
     Create vectorized LIBERO environments with a consistent return shape.
@@ -454,7 +437,7 @@ def create_libero_envs(
     gym_kwargs = dict(gym_kwargs or {})
     task_ids_filter = gym_kwargs.pop("task_ids", None)  # optional: limit to specific tasks
 
-    camera_names = parse_camera_names(camera_name)
+    camera_names = _parse_camera_names(camera_name)
     suite_names = [s.strip() for s in str(task).split(",") if s.strip()]
     if not suite_names:
         raise ValueError("`task` must contain at least one LIBERO suite name.")
@@ -479,7 +462,6 @@ def create_libero_envs(
         # Probe once and reuse to avoid creating a temp env per task.
         cached_obs_space: spaces.Space | None = None
         cached_act_space: spaces.Space | None = None
-        cached_metadata: dict[str, Any] | None = None
 
         for tid in selected:
             fns = _make_env_fns(
@@ -493,14 +475,12 @@ def create_libero_envs(
                 gym_kwargs=gym_kwargs,
                 control_mode=control_mode,
                 camera_name_mapping=camera_name_mapping,
-                is_libero_plus=is_libero_plus,
             )
             if is_async:
-                lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
+                lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space)
                 if cached_obs_space is None:
                     cached_obs_space = lazy.observation_space
                     cached_act_space = lazy.action_space
-                    cached_metadata = lazy.metadata
                 out[suite_name][tid] = lazy
             else:
                 out[suite_name][tid] = env_cls(fns)

src/lerobot/envs/metaworld.py

@@ -311,7 +311,6 @@ def create_metaworld_envs(
     is_async = env_cls is gym.vector.AsyncVectorEnv
     cached_obs_space = None
     cached_act_space = None
-    cached_metadata = None
     out: dict[str, dict[int, Any]] = defaultdict(dict)
 
     for group in task_groups:
@@ -325,11 +324,10 @@ def create_metaworld_envs(
             fns = [(lambda tn=task_name: MetaworldEnv(task=tn, **gym_kwargs)) for _ in range(n_envs)]
 
             if is_async:
-                lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
+                lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space)
                 if cached_obs_space is None:
                     cached_obs_space = lazy.observation_space
                     cached_act_space = lazy.action_space
-                    cached_metadata = lazy.metadata
                 out[group][tid] = lazy
             else:
                 out[group][tid] = env_cls(fns)

src/lerobot/envs/robocasa.py

@@ -1,425 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from __future__ import annotations
-
-import logging
-from collections import defaultdict
-from collections.abc import Callable, Sequence
-from functools import partial
-from typing import Any
-
-import gymnasium as gym
-import numpy as np
-from gymnasium import spaces
-
-from lerobot.types import RobotObservation
-
-from .utils import _LazyAsyncVectorEnv, parse_camera_names
-
-logger = logging.getLogger(__name__)
-
-# Dimensions for the flat action/state vectors used by the LeRobot wrapper.
-# These correspond to the PandaOmron robot in RoboCasa365.
-OBS_STATE_DIM = 16  # base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2)
-ACTION_DIM = 12  # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
-ACTION_LOW = -1.0
-ACTION_HIGH = 1.0
-
-# Default PandaOmron cameras. We surface these raw names directly as
-# `observation.images.<name>` so the LeRobot dataset/policy keys match
-# RoboCasa's native convention (no implicit renaming).
-DEFAULT_CAMERAS = [
-    "robot0_agentview_left",
-    "robot0_eye_in_hand",
-    "robot0_agentview_right",
-]
-
-# Object-mesh registries to sample from. RoboCasa's upstream default is
-# ("objaverse", "lightwheel"), but the objaverse pack is huge (~30GB) and
-# most users — including our CI image — only download the lightwheel pack
-# (`--type objs_lw` in `download_kitchen_assets`). When a sampled object
-# category has zero candidates in every registry, robocasa crashes with
-# `ValueError: Probabilities contain NaN` (0/0 divide in the probability
-# normalization). Restricting to registries that are actually on disk
-# avoids the NaN and matches what the asset download provides.
-DEFAULT_OBJ_REGISTRIES: tuple[str, ...] = ("lightwheel",)
-
-# Task-group shortcuts accepted as `--env.task`. When the user passes one of
-# these names, we expand it to the upstream RoboCasa task list and auto-set
-# the dataset split. Individual task names (optionally comma-separated) still
-# take precedence; this only triggers on an exact group-name match.
-_TASK_GROUP_SPLITS = {
-    "atomic_seen": "target",
-    "composite_seen": "target",
-    "composite_unseen": "target",
-    "pretrain50": "pretrain",
-    "pretrain100": "pretrain",
-    "pretrain200": "pretrain",
-    "pretrain300": "pretrain",
-}
-
-
-def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
-    """Resolve a `--env.task` value to (task_names, split_override).
-
-    If `task` is a known task-group name (e.g. `atomic_seen`, `pretrain100`),
-    expand it via `robocasa.utils.dataset_registry.{TARGET,PRETRAINING}_TASKS`
-    and return the matching split. Otherwise treat `task` as a single task or
-    comma-separated list and leave the split untouched (None).
-    """
-    key = task.strip()
-    if key in _TASK_GROUP_SPLITS:
-        from robocasa.utils.dataset_registry import PRETRAINING_TASKS, TARGET_TASKS
-
-        combined = {**TARGET_TASKS, **PRETRAINING_TASKS}
-        if key not in combined:
-            raise ValueError(
-                f"Task group '{key}' is not available in this version of robocasa. "
-                f"Known groups: {sorted(combined.keys())}."
-            )
-        return list(combined[key]), _TASK_GROUP_SPLITS[key]
-
-    names = [t.strip() for t in task.split(",") if t.strip()]
-    if not names:
-        raise ValueError("`task` must contain at least one RoboCasa task name.")
-    return names, None
-
-
-def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
-    """Split a flat (12,) action vector into a RoboCasa action dict.
-
-    Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
-    """
-    return {
-        "action.base_motion": flat_action[0:4],
-        "action.control_mode": flat_action[4:5],
-        "action.end_effector_position": flat_action[5:8],
-        "action.end_effector_rotation": flat_action[8:11],
-        "action.gripper_close": flat_action[11:12],
-    }
-
-
-class RoboCasaEnv(gym.Env):
-    """LeRobot gym.Env wrapper for RoboCasa365 kitchen environments.
-
-    Wraps RoboCasaGymEnv from the robocasa package and converts its
-    dict-based observations and actions into the flat arrays LeRobot expects.
-    Raw RoboCasa camera names are preserved verbatim under `pixels/<cam>`.
-    """
-
-    metadata = {"render_modes": ["rgb_array"], "render_fps": 20}
-
-    def __init__(
-        self,
-        task: str,
-        camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS),
-        obs_type: str = "pixels_agent_pos",
-        render_mode: str = "rgb_array",
-        observation_width: int = 256,
-        observation_height: int = 256,
-        visualization_width: int = 512,
-        visualization_height: int = 512,
-        split: str | None = None,
-        episode_length: int | None = None,
-        obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
-        episode_index: int = 0,
-    ):
-        super().__init__()
-        self.task = task
-        self.obs_type = obs_type
-        self.render_mode = render_mode
-        self.observation_width = observation_width
-        self.observation_height = observation_height
-        self.visualization_width = visualization_width
-        self.visualization_height = visualization_height
-        self.split = split
-        self.obj_registries = tuple(obj_registries)
-        # Per-worker index (0..n_envs-1) used to spread the user-provided
-        # seed across factories so each sub-env explores a distinct layout
-        # even when the same seed is passed to `reset()`.
-        self.episode_index = int(episode_index)
-
-        self.camera_name = parse_camera_names(camera_name)
-
-        self._max_episode_steps = episode_length if episode_length is not None else 1000
-
-        # Deferred — created on first reset() inside the worker subprocess
-        # to avoid inheriting stale GPU/EGL contexts across fork().
-        self._env: Any = None
-        self.task_description = ""
-
-        images = {
-            cam: spaces.Box(
-                low=0,
-                high=255,
-                shape=(self.observation_height, self.observation_width, 3),
-                dtype=np.uint8,
-            )
-            for cam in self.camera_name
-        }
-
-        if self.obs_type == "pixels":
-            self.observation_space = spaces.Dict({"pixels": spaces.Dict(images)})
-        elif self.obs_type == "pixels_agent_pos":
-            self.observation_space = spaces.Dict(
-                {
-                    "pixels": spaces.Dict(images),
-                    "agent_pos": spaces.Box(
-                        low=-np.inf,
-                        high=np.inf,
-                        shape=(OBS_STATE_DIM,),
-                        dtype=np.float32,
-                    ),
-                }
-            )
-        else:
-            raise ValueError(f"Unsupported obs_type '{self.obs_type}'. Use 'pixels' or 'pixels_agent_pos'.")
-
-        self.action_space = spaces.Box(
-            low=ACTION_LOW,
-            high=ACTION_HIGH,
-            shape=(ACTION_DIM,),
-            dtype=np.float32,
-        )
-
-    def _ensure_env(self) -> None:
-        """Create the underlying RoboCasaGymEnv on first use.
-
-        Called inside the worker subprocess after fork(), so each worker gets
-        its own clean rendering context rather than inheriting a stale one from
-        the parent process (which causes crashes with AsyncVectorEnv).
-        """
-        if self._env is not None:
-            return
-        from robocasa.wrappers.gym_wrapper import RoboCasaGymEnv
-
-        # RoboCasaGymEnv defaults split="test", which create_env rejects
-        # (only None/"all"/"pretrain"/"target" are valid). Always pass a
-        # valid value so we don't hit that default. Extra kwargs are
-        # forwarded to the underlying kitchen env via create_env/robosuite.make.
-        self._env = RoboCasaGymEnv(
-            env_name=self.task,
-            camera_widths=self.observation_width,
-            camera_heights=self.observation_height,
-            split=self.split if self.split is not None else "all",
-            obj_registries=self.obj_registries,
-        )
-
-        ep_meta = self._env.env.get_ep_meta()
-        self.task_description = ep_meta.get("lang", self.task)
-
-    def _format_raw_obs(self, raw_obs: dict) -> RobotObservation:
-        """Convert RoboCasaGymEnv observation dict to LeRobot format."""
-        # RoboCasaGymEnv emits camera frames under "video.<cam>".
-        images = {cam: raw_obs[f"video.{cam}"] for cam in self.camera_name if f"video.{cam}" in raw_obs}
-
-        if self.obs_type == "pixels":
-            return {"pixels": images}
-
-        # `state.*` keys come from PandaOmronKeyConverter inside the wrapper.
-        agent_pos = np.concatenate(
-            [
-                raw_obs.get("state.base_position", np.zeros(3)),
-                raw_obs.get("state.base_rotation", np.zeros(4)),
-                raw_obs.get("state.end_effector_position_relative", np.zeros(3)),
-                raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
-                raw_obs.get("state.gripper_qpos", np.zeros(2)),
-            ],
-            axis=-1,
-        ).astype(np.float32)
-
-        return {"pixels": images, "agent_pos": agent_pos}
-
-    def render(self) -> np.ndarray:
-        self._ensure_env()
-        assert self._env is not None
-        return self._env.render()
-
-    def reset(self, seed=None, **kwargs):
-        self._ensure_env()
-        assert self._env is not None
-        super().reset(seed=seed)
-        # Spread the seed across workers so n_envs factories don't all
-        # roll the same scene. With an explicit user seed we shift it by
-        # episode_index; with no seed we fall back to episode_index so
-        # each worker is still distinct rather than inheriting the same
-        # global RNG state.
-        worker_seed = seed + self.episode_index if seed is not None else self.episode_index
-        raw_obs, info = self._env.reset(seed=worker_seed)
-
-        ep_meta = self._env.env.get_ep_meta()
-        self.task_description = ep_meta.get("lang", self.task)
-
-        observation = self._format_raw_obs(raw_obs)
-        info = {"is_success": False}
-        return observation, info
-
-    def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]:
-        self._ensure_env()
-        assert self._env is not None
-        if action.ndim != 1:
-            raise ValueError(
-                f"Expected action to be 1-D (shape (action_dim,)), "
-                f"but got shape {action.shape} with ndim={action.ndim}"
-            )
-
-        action_dict = convert_action(action)
-        raw_obs, reward, done, truncated, info = self._env.step(action_dict)
-
-        is_success = bool(info.get("success", False))
-        terminated = done or is_success
-        info.update({"task": self.task, "done": done, "is_success": is_success})
-
-        observation = self._format_raw_obs(raw_obs)
-        if terminated:
-            info["final_info"] = {
-                "task": self.task,
-                "done": bool(done),
-                "is_success": bool(is_success),
-            }
-            self.reset()
-
-        return observation, reward, terminated, truncated, info
-
-    def close(self):
-        if self._env is not None:
-            self._env.close()
-
-
-def _make_env_fns(
-    *,
-    task: str,
-    n_envs: int,
-    camera_names: list[str],
-    obs_type: str,
-    render_mode: str,
-    observation_width: int,
-    observation_height: int,
-    visualization_width: int,
-    visualization_height: int,
-    split: str | None,
-    episode_length: int | None,
-    obj_registries: Sequence[str],
-) -> list[Callable[[], RoboCasaEnv]]:
-    """Build n_envs factory callables for a single task.
-
-    Each factory carries a distinct ``episode_index`` (``0..n_envs-1``) so
-    ``RoboCasaEnv.reset()`` can derive a per-worker seed series from the
-    user-provided seed.
-    """
-
-    def _make_env(episode_index: int) -> RoboCasaEnv:
-        return RoboCasaEnv(
-            task=task,
-            camera_name=camera_names,
-            obs_type=obs_type,
-            render_mode=render_mode,
-            observation_width=observation_width,
-            observation_height=observation_height,
-            visualization_width=visualization_width,
-            visualization_height=visualization_height,
-            split=split,
-            episode_length=episode_length,
-            obj_registries=obj_registries,
-            episode_index=episode_index,
-        )
-
-    return [partial(_make_env, i) for i in range(n_envs)]
-
-
-def create_robocasa_envs(
-    task: str,
-    n_envs: int,
-    gym_kwargs: dict[str, Any] | None = None,
-    camera_name: str | Sequence[str] = ",".join(DEFAULT_CAMERAS),
-    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
-    episode_length: int | None = None,
-    obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
-) -> dict[str, dict[int, Any]]:
-    """Create vectorized RoboCasa365 environments with a consistent return shape.
-
-    Returns:
-        dict[task_name][task_id] -> vec_env (env_cls([...]) with exactly n_envs factories)
-
-    `task` can be:
-      - a single task name (e.g. `CloseFridge`)
-      - a comma-separated list of task names (e.g. `CloseFridge,PickPlaceCoffee`)
-      - a benchmark-group shortcut (`atomic_seen`, `composite_seen`,
-        `composite_unseen`, `pretrain50`, `pretrain100`, `pretrain200`,
-        `pretrain300`), which auto-expands to the upstream task list and
-        auto-sets the dataset `split` ("target" or "pretrain").
-    """
-    if env_cls is None or not callable(env_cls):
-        raise ValueError("env_cls must be a callable that wraps a list of environment factory callables.")
-    if not isinstance(n_envs, int) or n_envs <= 0:
-        raise ValueError(f"n_envs must be a positive int; got {n_envs}.")
-
-    gym_kwargs = dict(gym_kwargs or {})
-    obs_type = gym_kwargs.pop("obs_type", "pixels_agent_pos")
-    render_mode = gym_kwargs.pop("render_mode", "rgb_array")
-    observation_width = gym_kwargs.pop("observation_width", 256)
-    observation_height = gym_kwargs.pop("observation_height", 256)
-    visualization_width = gym_kwargs.pop("visualization_width", 512)
-    visualization_height = gym_kwargs.pop("visualization_height", 512)
-    split = gym_kwargs.pop("split", None)
-
-    camera_names = parse_camera_names(camera_name)
-    task_names, group_split = _resolve_tasks(str(task))
-    if group_split is not None and split is None:
-        split = group_split
-
-    logger.info(
-        "Creating RoboCasa envs | tasks=%s | split=%s | n_envs(per task)=%d",
-        task_names,
-        split,
-        n_envs,
-    )
-
-    is_async = env_cls is gym.vector.AsyncVectorEnv
-
-    cached_obs_space: spaces.Space | None = None
-    cached_act_space: spaces.Space | None = None
-    cached_metadata: dict[str, Any] | None = None
-    out: dict[str, dict[int, Any]] = defaultdict(dict)
-
-    for task_name in task_names:
-        fns = _make_env_fns(
-            task=task_name,
-            n_envs=n_envs,
-            camera_names=camera_names,
-            obs_type=obs_type,
-            render_mode=render_mode,
-            observation_width=observation_width,
-            observation_height=observation_height,
-            visualization_width=visualization_width,
-            visualization_height=visualization_height,
-            split=split,
-            episode_length=episode_length,
-            obj_registries=obj_registries,
-        )
-
-        if is_async:
-            lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
-            if cached_obs_space is None:
-                cached_obs_space = lazy.observation_space
-                cached_act_space = lazy.action_space
-                cached_metadata = lazy.metadata
-            out[task_name][0] = lazy
-        else:
-            out[task_name][0] = env_cls(fns)
-        logger.info("Built vec env | task=%s | n_envs=%d", task_name, n_envs)
-
-    return {name: dict(task_map) for name, task_map in out.items()}

src/lerobot/envs/robomme.py

@@ -1,245 +0,0 @@
-"""RoboMME environment wrapper for LeRobot evaluation.
-
-Wraps the RoboMME ``BenchmarkEnvBuilder`` into a Gymnasium-compatible
-``VectorEnv`` suitable for ``lerobot_eval``.
-
-RoboMME tasks:
-  Counting:    BinFill, PickXtimes, SwingXtimes, StopCube
-  Permanence:  VideoUnmask, VideoUnmaskSwap, ButtonUnmask, ButtonUnmaskSwap
-  Reference:   PickHighlight, VideoRepick, VideoPlaceButton, VideoPlaceOrder
-  Imitation:   MoveCube, InsertPeg, PatternLock, RouteStick
-
-Dataset: lerobot/robomme (LeRobot v3.0, 1,600 episodes)
-Install: see docker/Dockerfile.benchmark.robomme  (Linux only — mani-skill vs numpy pin conflict)
-Benchmark: https://github.com/RoboMME/robomme_benchmark
-"""
-
-from __future__ import annotations
-
-from collections.abc import Callable, Sequence
-from functools import partial
-from typing import Any
-
-import gymnasium as gym
-import numpy as np
-from gymnasium import spaces
-
-from .utils import _LazyAsyncVectorEnv
-
-ROBOMME_TASKS = [
-    "BinFill",
-    "PickXtimes",
-    "SwingXtimes",
-    "StopCube",
-    "VideoUnmask",
-    "VideoUnmaskSwap",
-    "ButtonUnmask",
-    "ButtonUnmaskSwap",
-    "PickHighlight",
-    "VideoRepick",
-    "VideoPlaceButton",
-    "VideoPlaceOrder",
-    "MoveCube",
-    "InsertPeg",
-    "PatternLock",
-    "RouteStick",
-]
-
-
-class RoboMMEGymEnv(gym.Env):
-    """Thin Gymnasium wrapper around a single RoboMME episode env."""
-
-    metadata = {"render_modes": ["rgb_array"], "render_fps": 10}
-
-    def __init__(
-        self,
-        task: str = "PickXtimes",
-        action_space_type: str = "joint_angle",
-        dataset: str = "test",
-        episode_idx: int = 0,
-        max_steps: int = 300,
-    ):
-        super().__init__()
-        from robomme.env_record_wrapper import BenchmarkEnvBuilder
-
-        self._task = task
-        self._action_space_type = action_space_type
-        self._dataset = dataset
-        self._episode_idx = episode_idx
-        self._max_steps = max_steps
-        self._max_episode_steps = max_steps
-
-        self._builder = BenchmarkEnvBuilder(
-            env_id=task,
-            dataset=dataset,
-            action_space=action_space_type,
-            gui_render=False,
-            max_steps=max_steps,
-        )
-        self._env = None
-        self._last_raw_obs: dict | None = None
-
-        action_dim = 8 if action_space_type == "joint_angle" else 7
-        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(action_dim,), dtype=np.float32)
-        # `pixels` must be a nested Dict so `preprocess_observation()` in
-        # envs/utils.py picks it up and maps each camera to
-        # `observation.images.<cam>`. A flat layout (`pixels/image`,
-        # `pixels/wrist_image`) silently drops every image from the batch.
-        self.observation_space = spaces.Dict(
-            {
-                "pixels": spaces.Dict(
-                    {
-                        "image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8),
-                        "wrist_image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8),
-                    }
-                ),
-                "agent_pos": spaces.Box(-np.inf, np.inf, shape=(8,), dtype=np.float32),
-            }
-        )
-
-    def reset(self, *, seed=None, options=None):
-        super().reset(seed=seed)
-        self._env = self._builder.make_env_for_episode(
-            episode_idx=self._episode_idx,
-            max_steps=self._max_steps,
-        )
-        obs, info = self._env.reset()
-        self._last_raw_obs = obs
-        return self._convert_obs(obs), self._convert_info(info)
-
-    def step(self, action):
-        obs, reward, terminated, truncated, info = self._env.step(action)
-        self._last_raw_obs = obs
-
-        terminated_bool = bool(terminated.item()) if hasattr(terminated, "item") else bool(terminated)
-        truncated_bool = bool(truncated.item()) if hasattr(truncated, "item") else bool(truncated)
-
-        status = info.get("status", "ongoing")
-        is_success = status == "success"
-        conv_info = self._convert_info(info)
-        conv_info["is_success"] = is_success
-
-        return self._convert_obs(obs), float(reward), terminated_bool, truncated_bool, conv_info
-
-    def render(self) -> np.ndarray | None:
-        """Return the front camera image from the last observation for video recording."""
-        if self._last_raw_obs is None:
-            return np.zeros((256, 256, 3), dtype=np.uint8)
-        front = self._last_raw_obs.get("front_rgb_list")
-        if front is None:
-            return np.zeros((256, 256, 3), dtype=np.uint8)
-        frame = front[-1] if isinstance(front, list) else front
-        return np.asarray(frame, dtype=np.uint8)
-
-    def _convert_obs(self, obs: dict) -> dict:
-        front_rgb = (
-            obs["front_rgb_list"][-1] if isinstance(obs["front_rgb_list"], list) else obs["front_rgb_list"]
-        )
-        wrist_rgb = (
-            obs["wrist_rgb_list"][-1] if isinstance(obs["wrist_rgb_list"], list) else obs["wrist_rgb_list"]
-        )
-        joint_state = (
-            obs["joint_state_list"][-1]
-            if isinstance(obs["joint_state_list"], list)
-            else obs["joint_state_list"]
-        )
-        gripper_state = (
-            obs["gripper_state_list"][-1]
-            if isinstance(obs["gripper_state_list"], list)
-            else obs["gripper_state_list"]
-        )
-
-        front_rgb = np.asarray(front_rgb, dtype=np.uint8)
-        wrist_rgb = np.asarray(wrist_rgb, dtype=np.uint8)
-        joint = np.asarray(joint_state, dtype=np.float32).flatten()[:7]
-        gripper = np.asarray(gripper_state, dtype=np.float32).flatten()[:1]
-        state = np.concatenate([joint, gripper])
-
-        return {
-            "pixels": {"image": front_rgb, "wrist_image": wrist_rgb},
-            "agent_pos": state,
-        }
-
-    def _convert_info(self, info: dict) -> dict:
-        return {
-            "status": info.get("status", "ongoing"),
-            "task_goal": info.get("task_goal", ""),
-        }
-
-
-def _make_env_fns(
-    *,
-    task: str,
-    n_envs: int,
-    action_space_type: str,
-    dataset: str,
-    episode_length: int,
-    task_id: int,
-) -> list[Callable[[], RoboMMEGymEnv]]:
-    """Build n_envs factory callables for one RoboMME task id."""
-
-    def _make_one(episode_index: int) -> RoboMMEGymEnv:
-        return RoboMMEGymEnv(
-            task=task,
-            action_space_type=action_space_type,
-            dataset=dataset,
-            episode_idx=episode_index,
-            max_steps=episode_length,
-        )
-
-    return [partial(_make_one, task_id + i) for i in range(n_envs)]
-
-
-def create_robomme_envs(
-    task: str,
-    n_envs: int = 1,
-    action_space_type: str = "joint_angle",
-    dataset: str = "test",
-    episode_length: int = 300,
-    task_ids: list[int] | None = None,
-    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
-) -> dict[str, dict[int, gym.vector.VectorEnv]]:
-    """Create vectorized RoboMME environments for evaluation.
-
-    `task` may be a single RoboMME task name (e.g. "PickXtimes") or a
-    comma-separated list (e.g. "PickXtimes,BinFill,StopCube"). Each task
-    becomes its own suite in the returned mapping.
-
-    Returns {suite_name: {task_id: VectorEnv}} matching lerobot's expected format.
-    """
-    if env_cls is None or not callable(env_cls):
-        raise ValueError("env_cls must be a callable that wraps a list of env factory callables.")
-    if not isinstance(n_envs, int) or n_envs <= 0:
-        raise ValueError(f"n_envs must be a positive int; got {n_envs}.")
-
-    if task_ids is None:
-        task_ids = [0]
-
-    task_names = [t.strip() for t in task.split(",") if t.strip()]
-    is_async = env_cls is gym.vector.AsyncVectorEnv
-    cached_obs_space: spaces.Space | None = None
-    cached_act_space: spaces.Space | None = None
-    cached_metadata: dict[str, Any] | None = None
-    out: dict[str, dict[int, gym.vector.VectorEnv]] = {}
-    for task_name in task_names:
-        envs_by_task: dict[int, gym.vector.VectorEnv] = {}
-        for task_id in task_ids:
-            fns = _make_env_fns(
-                task=task_name,
-                n_envs=n_envs,
-                action_space_type=action_space_type,
-                dataset=dataset,
-                episode_length=episode_length,
-                task_id=task_id,
-            )
-            if is_async:
-                lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
-                if cached_obs_space is None:
-                    cached_obs_space = lazy.observation_space
-                    cached_act_space = lazy.action_space
-                    cached_metadata = lazy.metadata
-                envs_by_task[task_id] = lazy
-            else:
-                envs_by_task[task_id] = env_cls(fns)
-        out[task_name] = envs_by_task
-    return out

src/lerobot/envs/robotwin.py

@@ -1,488 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from __future__ import annotations
-
-import importlib
-import logging
-from collections import defaultdict
-from collections.abc import Callable, Sequence
-from functools import partial
-from typing import Any
-
-import gymnasium as gym
-import numpy as np
-import torch
-from gymnasium import spaces
-
-from lerobot.types import RobotObservation
-
-from .utils import _LazyAsyncVectorEnv
-
-logger = logging.getLogger(__name__)
-
-# Camera names as used by RoboTwin 2.0. The wrapper appends "_rgb" when looking
-# up keys in get_obs() output (e.g. "head_camera" → "head_camera_rgb").
-ROBOTWIN_CAMERA_NAMES: tuple[str, ...] = (
-    "head_camera",
-    "left_camera",
-    "right_camera",
-)
-
-ACTION_DIM = 14  # 7 DOF × 2 arms
-ACTION_LOW = -1.0
-ACTION_HIGH = 1.0
-DEFAULT_EPISODE_LENGTH = 300
-# D435 dims from task_config/_camera_config.yml (what demo_clean.yml selects).
-DEFAULT_CAMERA_H = 240
-DEFAULT_CAMERA_W = 320
-
-# Task list from RoboTwin 2.0's `envs/` directory — mirrors upstream exactly
-# (50 tasks as of main; earlier revisions had 60 with a different split).
-# Keep this in sync with:
-#   gh api /repos/RoboTwin-Platform/RoboTwin/contents/envs --paginate \
-#     | jq -r '.[].name' | grep -E '\.py$' | grep -v '^_' | sed 's/\.py$//'
-ROBOTWIN_TASKS: tuple[str, ...] = (
-    "adjust_bottle",
-    "beat_block_hammer",
-    "blocks_ranking_rgb",
-    "blocks_ranking_size",
-    "click_alarmclock",
-    "click_bell",
-    "dump_bin_bigbin",
-    "grab_roller",
-    "handover_block",
-    "handover_mic",
-    "hanging_mug",
-    "lift_pot",
-    "move_can_pot",
-    "move_pillbottle_pad",
-    "move_playingcard_away",
-    "move_stapler_pad",
-    "open_laptop",
-    "open_microwave",
-    "pick_diverse_bottles",
-    "pick_dual_bottles",
-    "place_a2b_left",
-    "place_a2b_right",
-    "place_bread_basket",
-    "place_bread_skillet",
-    "place_burger_fries",
-    "place_can_basket",
-    "place_cans_plasticbox",
-    "place_container_plate",
-    "place_dual_shoes",
-    "place_empty_cup",
-    "place_fan",
-    "place_mouse_pad",
-    "place_object_basket",
-    "place_object_scale",
-    "place_object_stand",
-    "place_phone_stand",
-    "place_shoe",
-    "press_stapler",
-    "put_bottles_dustbin",
-    "put_object_cabinet",
-    "rotate_qrcode",
-    "scan_object",
-    "shake_bottle",
-    "shake_bottle_horizontally",
-    "stack_blocks_three",
-    "stack_blocks_two",
-    "stack_bowls_three",
-    "stack_bowls_two",
-    "stamp_seal",
-    "turn_switch",
-)
-
-
-_ROBOTWIN_SETUP_CACHE: dict[str, dict[str, Any]] = {}
-
-
-def _load_robotwin_setup_kwargs(task_name: str) -> dict[str, Any]:
-    """Build the kwargs dict RoboTwin's setup_demo expects.
-
-    Mirrors the config loading done by RoboTwin's ``script/eval_policy.py``:
-    reads ``task_config/demo_clean.yml``, resolves the embodiment file from
-    ``_embodiment_config.yml``, loads the robot's own ``config.yml``, and
-    reads camera dimensions from ``_camera_config.yml``.
-
-    Uses ``aloha-agilex`` single-robot dual-arm by default (the only embodiment
-    used by beat_block_hammer and most smoke-test tasks).
-    """
-    if task_name in _ROBOTWIN_SETUP_CACHE:
-        return dict(_ROBOTWIN_SETUP_CACHE[task_name])
-
-    import os
-
-    import yaml  # type: ignore[import-untyped]
-    from envs import CONFIGS_PATH  # type: ignore[import-not-found]
-
-    task_config = "demo_clean"
-    with open(os.path.join(CONFIGS_PATH, f"{task_config}.yml"), encoding="utf-8") as f:
-        args = yaml.safe_load(f)
-
-    # Resolve embodiment — demo_clean.yml uses [aloha-agilex] (dual-arm single robot)
-    with open(os.path.join(CONFIGS_PATH, "_embodiment_config.yml"), encoding="utf-8") as f:
-        embodiment_types = yaml.safe_load(f)
-    embodiment = args.get("embodiment", ["aloha-agilex"])
-    if len(embodiment) == 1:
-        robot_file = embodiment_types[embodiment[0]]["file_path"]
-        args["left_robot_file"] = robot_file
-        args["right_robot_file"] = robot_file
-        args["dual_arm_embodied"] = True
-    elif len(embodiment) == 3:
-        args["left_robot_file"] = embodiment_types[embodiment[0]]["file_path"]
-        args["right_robot_file"] = embodiment_types[embodiment[1]]["file_path"]
-        args["embodiment_dis"] = embodiment[2]
-        args["dual_arm_embodied"] = False
-    else:
-        raise ValueError(f"embodiment must have 1 or 3 items, got {len(embodiment)}")
-
-    with open(os.path.join(args["left_robot_file"], "config.yml"), encoding="utf-8") as f:
-        args["left_embodiment_config"] = yaml.safe_load(f)
-    with open(os.path.join(args["right_robot_file"], "config.yml"), encoding="utf-8") as f:
-        args["right_embodiment_config"] = yaml.safe_load(f)
-
-    # Camera dimensions
-    with open(os.path.join(CONFIGS_PATH, "_camera_config.yml"), encoding="utf-8") as f:
-        camera_config = yaml.safe_load(f)
-    head_cam = args["camera"]["head_camera_type"]
-    args["head_camera_h"] = camera_config[head_cam]["h"]
-    args["head_camera_w"] = camera_config[head_cam]["w"]
-
-    # Headless overrides
-    args["render_freq"] = 0
-    args["task_name"] = task_name
-    args["task_config"] = task_config
-
-    _ROBOTWIN_SETUP_CACHE[task_name] = args
-    return dict(args)
-
-
-def _load_robotwin_task(task_name: str) -> type:
-    """Dynamically import and return a RoboTwin 2.0 task class.
-
-    RoboTwin tasks live in ``envs/<task_name>.py`` relative to the repository
-    root and are expected to be on ``sys.path`` after installation.
-    """
-    try:
-        module = importlib.import_module(f"envs.{task_name}")
-    except ModuleNotFoundError as e:
-        raise ModuleNotFoundError(
-            f"Could not import RoboTwin task '{task_name}'. "
-            "Ensure RoboTwin 2.0 is installed and its 'envs/' directory is on PYTHONPATH. "
-            "See the RoboTwin installation guide: https://robotwin-platform.github.io/doc/usage/robotwin-install.html"
-        ) from e
-    task_cls = getattr(module, task_name, None)
-    if task_cls is None:
-        raise AttributeError(f"Task class '{task_name}' not found in envs/{task_name}.py")
-    return task_cls
-
-
-class RoboTwinEnv(gym.Env):
-    """Gymnasium wrapper around a single RoboTwin 2.0 task.
-
-    RoboTwin uses a custom SAPIEN-based API (``setup_demo`` / ``get_obs`` /
-    ``take_action`` / ``check_success``) rather than the standard gym interface.
-    This class bridges that API to Gymnasium so that ``lerobot-eval`` can drive
-    RoboTwin exactly like LIBERO or Meta-World.
-
-    The underlying SAPIEN environment is created lazily on the first ``reset()``
-    call *inside the worker process*.  This is required for
-    ``gym.vector.AsyncVectorEnv`` compatibility: SAPIEN allocates EGL/GPU
-    contexts that must not be forked from the parent process.
-
-    Observations
-    ------------
-    The ``pixels`` dict uses the raw RoboTwin camera names as keys (e.g.
-    ``"head_camera"``, ``"left_camera"``). ``preprocess_observation`` in
-    ``envs/utils.py`` then converts these to ``observation.images.<cam>``.
-
-    Actions
-    -------
-    14-dim float32 array in ``[-1, 1]`` (joint-space, 7 DOF per arm).
-
-    Autograd
-    --------
-    ``setup_demo`` and ``take_action`` drive CuRobo's Newton trajectory
-    optimizer, which calls ``cost.backward()`` internally. lerobot_eval wraps
-    the rollout in ``torch.no_grad()``, so both call sites re-enable grad.
-    """
-
-    metadata = {"render_modes": ["rgb_array"], "render_fps": 25}
-
-    def __init__(
-        self,
-        task_name: str,
-        episode_index: int = 0,
-        n_envs: int = 1,
-        camera_names: Sequence[str] = ROBOTWIN_CAMERA_NAMES,
-        observation_height: int | None = None,
-        observation_width: int | None = None,
-        episode_length: int = DEFAULT_EPISODE_LENGTH,
-        render_mode: str = "rgb_array",
-    ):
-        super().__init__()
-        self.task_name = task_name
-        self.task = task_name  # used by add_envs_task() in utils.py
-        self.task_description = task_name.replace("_", " ")
-        self.episode_index = episode_index
-        self._reset_stride = n_envs
-        self.camera_names = list(camera_names)
-        # Default to D435 dims (the camera type baked into task_config/demo_clean.yml).
-        # The YAML-driven lookup is deferred to reset() so construction doesn't
-        # import RoboTwin's `envs` module — fast-tests run without RoboTwin installed.
-        self.observation_height = observation_height or DEFAULT_CAMERA_H
-        self.observation_width = observation_width or DEFAULT_CAMERA_W
-        self.episode_length = episode_length
-        self._max_episode_steps = episode_length  # lerobot_eval.rollout reads this
-        self.render_mode = render_mode
-
-        self._env: Any | None = None  # deferred — created on first reset() inside worker
-        self._step_count: int = 0
-        self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
-
-        image_spaces = {
-            cam: spaces.Box(
-                low=0,
-                high=255,
-                shape=(self.observation_height, self.observation_width, 3),
-                dtype=np.uint8,
-            )
-            for cam in self.camera_names
-        }
-        self.observation_space = spaces.Dict(
-            {
-                "pixels": spaces.Dict(image_spaces),
-                "agent_pos": spaces.Box(low=-np.inf, high=np.inf, shape=(ACTION_DIM,), dtype=np.float32),
-            }
-        )
-        self.action_space = spaces.Box(
-            low=ACTION_LOW, high=ACTION_HIGH, shape=(ACTION_DIM,), dtype=np.float32
-        )
-
-    def _ensure_env(self) -> None:
-        """Create the SAPIEN environment on first use.
-
-        Called inside the worker subprocess after fork(), so each worker gets
-        its own EGL/GPU context rather than inheriting a stale one from the
-        parent process (which causes crashes with AsyncVectorEnv).
-        """
-        if self._env is not None:
-            return
-        task_cls = _load_robotwin_task(self.task_name)
-        self._env = task_cls()
-
-    def _get_obs(self) -> RobotObservation:
-        assert self._env is not None, "_get_obs called before _ensure_env()"
-        raw = self._env.get_obs()
-        cameras_raw = raw.get("observation", {})
-
-        images: dict[str, np.ndarray] = {}
-        for cam in self.camera_names:
-            cam_data = cameras_raw.get(cam)
-            img = cam_data.get("rgb") if cam_data else None
-            if img is None:
-                images[cam] = self._black_frame
-                continue
-            img = np.asarray(img, dtype=np.uint8)
-            if img.ndim == 2:
-                img = np.stack([img, img, img], axis=-1)
-            elif img.shape[-1] != 3:
-                img = img[..., :3]
-            images[cam] = img
-
-        ja = raw.get("joint_action") or {}
-        vec = ja.get("vector")
-        if vec is not None:
-            arr = np.asarray(vec, dtype=np.float32).ravel()
-            joint_state = (
-                arr[:ACTION_DIM] if arr.size >= ACTION_DIM else np.zeros(ACTION_DIM, dtype=np.float32)
-            )
-        else:
-            joint_state = np.zeros(ACTION_DIM, dtype=np.float32)
-
-        return {"pixels": images, "agent_pos": joint_state}
-
-    def reset(self, seed: int | None = None, **kwargs) -> tuple[RobotObservation, dict]:
-        self._ensure_env()
-        super().reset(seed=seed)
-        assert self._env is not None  # set by _ensure_env() above
-
-        actual_seed = self.episode_index if seed is None else seed
-        setup_kwargs = _load_robotwin_setup_kwargs(self.task_name)
-        setup_kwargs.update(seed=actual_seed, is_test=True)
-        with torch.enable_grad():
-            self._env.setup_demo(**setup_kwargs)
-        self.episode_index += self._reset_stride
-        self._step_count = 0
-
-        obs = self._get_obs()
-        return obs, {"is_success": False, "task": self.task_name}
-
-    def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]:
-        assert self._env is not None, "step() called before reset()"
-        if action.ndim != 1 or action.shape[0] != ACTION_DIM:
-            raise ValueError(f"Expected 1-D action of shape ({ACTION_DIM},), got {action.shape}")
-
-        with torch.enable_grad():
-            if hasattr(self._env, "take_action"):
-                self._env.take_action(action)
-            else:
-                self._env.step(action)
-
-        self._step_count += 1
-
-        is_success = bool(getattr(self._env, "eval_success", False))
-        if not is_success and hasattr(self._env, "check_success"):
-            is_success = bool(self._env.check_success())
-
-        obs = self._get_obs()
-        reward = float(is_success)
-        terminated = is_success
-        truncated = self._step_count >= self.episode_length
-
-        info: dict[str, Any] = {
-            "task": self.task_name,
-            "is_success": is_success,
-            "step": self._step_count,
-        }
-        if terminated or truncated:
-            info["final_info"] = {
-                "task": self.task_name,
-                "is_success": is_success,
-            }
-            self.reset()
-
-        return obs, reward, terminated, truncated, info
-
-    def render(self) -> np.ndarray:
-        self._ensure_env()
-        obs = self._get_obs()
-        # Prefer head camera for rendering; fall back to first available.
-        if "head_camera" in obs["pixels"]:
-            return obs["pixels"]["head_camera"]
-        return next(iter(obs["pixels"].values()))
-
-    def close(self) -> None:
-        if self._env is not None:
-            if hasattr(self._env, "close_env"):
-                import contextlib
-
-                with contextlib.suppress(TypeError):
-                    self._env.close_env()
-            self._env = None
-
-
-# ---- Multi-task factory --------------------------------------------------------
-
-
-def _make_env_fns(
-    *,
-    task_name: str,
-    n_envs: int,
-    camera_names: list[str],
-    observation_height: int,
-    observation_width: int,
-    episode_length: int,
-) -> list[Callable[[], RoboTwinEnv]]:
-    """Return n_envs factory callables for a single task."""
-
-    def _make_one(episode_index: int) -> RoboTwinEnv:
-        return RoboTwinEnv(
-            task_name=task_name,
-            episode_index=episode_index,
-            n_envs=n_envs,
-            camera_names=camera_names,
-            observation_height=observation_height,
-            observation_width=observation_width,
-            episode_length=episode_length,
-        )
-
-    return [partial(_make_one, i) for i in range(n_envs)]
-
-
-def create_robotwin_envs(
-    task: str,
-    n_envs: int,
-    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
-    camera_names: Sequence[str] = ROBOTWIN_CAMERA_NAMES,
-    observation_height: int = DEFAULT_CAMERA_H,
-    observation_width: int = DEFAULT_CAMERA_W,
-    episode_length: int = DEFAULT_EPISODE_LENGTH,
-) -> dict[str, dict[int, Any]]:
-    """Create vectorized RoboTwin 2.0 environments.
-
-    Returns:
-        ``dict[task_name][0] -> VectorEnv`` — one entry per task, each wrapping
-        ``n_envs`` parallel rollouts.
-
-    Args:
-        task: Comma-separated list of task names (e.g. ``"beat_block_hammer"``
-            or ``"beat_block_hammer,click_bell"``).
-        n_envs: Number of parallel rollouts per task.
-        env_cls: Vector env constructor (e.g. ``gym.vector.AsyncVectorEnv``).
-        camera_names: Cameras to include in observations.
-        observation_height: Pixel height for all cameras.
-        observation_width: Pixel width for all cameras.
-        episode_length: Max steps before truncation.
-    """
-    if env_cls is None or not callable(env_cls):
-        raise ValueError("env_cls must be callable (e.g. gym.vector.AsyncVectorEnv).")
-    if not isinstance(n_envs, int) or n_envs <= 0:
-        raise ValueError(f"n_envs must be a positive int; got {n_envs}.")
-
-    task_names = [t.strip() for t in str(task).split(",") if t.strip()]
-    if not task_names:
-        raise ValueError("`task` must contain at least one RoboTwin task name.")
-
-    unknown = [t for t in task_names if t not in ROBOTWIN_TASKS]
-    if unknown:
-        raise ValueError(f"Unknown RoboTwin tasks: {unknown}. Available tasks: {sorted(ROBOTWIN_TASKS)}")
-
-    logger.info(
-        "Creating RoboTwin envs | tasks=%s | n_envs(per task)=%d",
-        task_names,
-        n_envs,
-    )
-
-    is_async = env_cls is gym.vector.AsyncVectorEnv
-    cached_obs_space: spaces.Space | None = None
-    cached_act_space: spaces.Space | None = None
-    cached_metadata: dict[str, Any] | None = None
-
-    out: dict[str, dict[int, Any]] = defaultdict(dict)
-    for task_name in task_names:
-        fns = _make_env_fns(
-            task_name=task_name,
-            n_envs=n_envs,
-            camera_names=list(camera_names),
-            observation_height=observation_height,
-            observation_width=observation_width,
-            episode_length=episode_length,
-        )
-        if is_async:
-            lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
-            if cached_obs_space is None:
-                cached_obs_space = lazy.observation_space
-                cached_act_space = lazy.action_space
-                cached_metadata = lazy.metadata
-            out[task_name][0] = lazy
-        else:
-            out[task_name][0] = env_cls(fns)
-        logger.info("Built vec env | task=%s | n_envs=%d", task_name, n_envs)
-
-    return {k: dict(v) for k, v in out.items()}

src/lerobot/envs/utils.py

@@ -34,25 +34,6 @@ from lerobot.utils.utils import get_channel_first_image_shape
 from .configs import EnvConfig
 
 
-def parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
-    """Normalize ``camera_name`` into a non-empty list of strings.
-
-    Accepts a comma-separated string (``"cam_a,cam_b"``) or a sequence of
-    strings (tuples/lists). Whitespace is stripped; empty entries are
-    dropped. Raises ``TypeError`` for unsupported input types and
-    ``ValueError`` when the normalized list is empty.
-    """
-    if isinstance(camera_name, str):
-        cams = [c.strip() for c in camera_name.split(",") if c.strip()]
-    elif isinstance(camera_name, (list | tuple)):
-        cams = [str(c).strip() for c in camera_name if str(c).strip()]
-    else:
-        raise TypeError(f"camera_name must be str or sequence[str], got {type(camera_name).__name__}")
-    if not cams:
-        raise ValueError("camera_name resolved to an empty list.")
-    return cams
-
-
 def _convert_nested_dict(d):
     result = {}
     for k, v in d.items():
@@ -172,20 +153,17 @@ class _LazyAsyncVectorEnv:
         env_fns: list[Callable],
         observation_space=None,
         action_space=None,
-        metadata=None,
     ):
         self._env_fns = env_fns
         self._env: gym.vector.AsyncVectorEnv | None = None
         self.num_envs = len(env_fns)
-        if observation_space is not None and action_space is not None and metadata is not None:
+        if observation_space is not None and action_space is not None:
             self.observation_space = observation_space
             self.action_space = action_space
-            self.metadata = metadata
         else:
             tmp = env_fns[0]()
             self.observation_space = tmp.observation_space
             self.action_space = tmp.action_space
-            self.metadata = tmp.metadata
             tmp.close()
         self.single_observation_space = self.observation_space
         self.single_action_space = self.action_space
@@ -194,10 +172,6 @@ class _LazyAsyncVectorEnv:
         if self._env is None:
             self._env = gym.vector.AsyncVectorEnv(self._env_fns, context="forkserver", shared_memory=True)
 
-    @property
-    def unwrapped(self):
-        return self
-
     def reset(self, **kwargs):
         self._ensure()
         return self._env.reset(**kwargs)

src/lerobot/envs/vlabench.py

@@ -1,589 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""VLABench environment wrapper for LeRobot.
-
-VLABench is a large-scale benchmark for language-conditioned robotic manipulation
-with long-horizon reasoning, built on MuJoCo/dm_control.
-
-- Paper: https://arxiv.org/abs/2412.18194
-- GitHub: https://github.com/OpenMOSS/VLABench
-- Website: https://vlabench.github.io
-"""
-
-from __future__ import annotations
-
-import contextlib
-import logging
-from collections import defaultdict
-from collections.abc import Callable, Sequence
-from typing import Any
-
-import cv2
-import gymnasium as gym
-import numpy as np
-from gymnasium import spaces
-from scipy.spatial.transform import Rotation
-
-from lerobot.types import RobotObservation
-
-from .utils import _LazyAsyncVectorEnv
-
-logger = logging.getLogger(__name__)
-
-ACTION_DIM = 7  # pos(3) + euler(3) + gripper(1)
-ACTION_LOW = np.array([-1.0, -1.0, -1.0, -1.0, -1.0, -1.0, 0.0], dtype=np.float32)
-ACTION_HIGH = np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], dtype=np.float32)
-
-# Default max episode steps per task type
-DEFAULT_MAX_EPISODE_STEPS = 500
-
-# VLABench task suites
-PRIMITIVE_TASKS = [
-    "select_fruit",
-    "select_toy",
-    "select_chemistry_tube",
-    "add_condiment",
-    "select_book",
-    "select_painting",
-    "select_drink",
-    "insert_flower",
-    "select_billiards",
-    "select_ingredient",
-    "select_mahjong",
-    "select_poker",
-    # Physical series
-    "density_qa",
-    "friction_qa",
-    "magnetism_qa",
-    "reflection_qa",
-    "simple_cuestick_usage",
-    "simple_seesaw_usage",
-    "sound_speed_qa",
-    "thermal_expansion_qa",
-    "weight_qa",
-]
-
-COMPOSITE_TASKS = [
-    "cluster_billiards",
-    "cluster_book",
-    "cluster_drink",
-    "cluster_toy",
-    "cook_dishes",
-    "cool_drink",
-    "find_unseen_object",
-    "get_coffee",
-    "hammer_nail",
-    "heat_food",
-    "make_juice",
-    "play_mahjong",
-    "play_math_game",
-    "play_poker",
-    "play_snooker",
-    "rearrange_book",
-    "rearrange_chemistry_tube",
-    "set_dining_table",
-    "set_study_table",
-    "store_food",
-    "take_chemistry_experiment",
-    "use_seesaw_complex",
-]
-
-SUITE_TASKS: dict[str, list[str]] = {
-    "primitive": PRIMITIVE_TASKS,
-    "composite": COMPOSITE_TASKS,
-}
-
-
-class VLABenchEnv(gym.Env):
-    """Gymnasium wrapper for VLABench environments.
-
-    Wraps the dm_control-based VLABench simulator behind a standard gym.Env interface.
-    Supports multiple cameras (front, second, wrist) and end-effector control.
-    """
-
-    metadata = {"render_modes": ["rgb_array"], "render_fps": 10}
-
-    def __init__(
-        self,
-        task: str = "select_fruit",
-        obs_type: str = "pixels_agent_pos",
-        render_mode: str = "rgb_array",
-        render_resolution: tuple[int, int] = (480, 480),
-        robot: str = "franka",
-        max_episode_steps: int = DEFAULT_MAX_EPISODE_STEPS,
-        action_mode: str = "eef",
-    ):
-        super().__init__()
-        self.task = task
-        self.obs_type = obs_type
-        self.render_mode = render_mode
-        self.render_resolution = render_resolution
-        self.robot = robot
-        self._max_episode_steps = max_episode_steps
-        self.action_mode = action_mode
-
-        # Deferred — created on first reset() inside worker subprocess to avoid
-        # inheriting stale GPU/EGL contexts when AsyncVectorEnv spawns workers.
-        # We never cache `env.physics`: dm_control exposes it as a weakref
-        # proxy that goes stale across resets (rebuilds the sim), so we always
-        # refetch it via `self._env.physics` at the call site.
-        self._env = None
-        self.task_description = ""  # populated on first reset
-        # Cached world-frame XYZ of the robot base link. The VLABench datasets
-        # log both `observation.state` positions and `actions` positions in
-        # robot-base frame (see VLABench/scripts/convert_to_lerobot.py which
-        # subtracts `robot_frame_pos` from ee_pos). The robot is attached at a
-        # fixed offset per task so this is safe to cache once per env build.
-        self._robot_base_xyz: np.ndarray | None = None
-
-        h, w = self.render_resolution
-
-        if self.obs_type == "state":
-            raise NotImplementedError(
-                "The 'state' observation type is not supported in VLABenchEnv. "
-                "Please use 'pixels' or 'pixels_agent_pos'."
-            )
-        elif self.obs_type == "pixels":
-            self.observation_space = spaces.Dict(
-                {
-                    "pixels": spaces.Dict(
-                        {
-                            "image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8),
-                            "second_image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8),
-                            "wrist_image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8),
-                        }
-                    ),
-                }
-            )
-        elif self.obs_type == "pixels_agent_pos":
-            self.observation_space = spaces.Dict(
-                {
-                    "pixels": spaces.Dict(
-                        {
-                            "image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8),
-                            "second_image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8),
-                            "wrist_image": spaces.Box(low=0, high=255, shape=(h, w, 3), dtype=np.uint8),
-                        }
-                    ),
-                    "agent_pos": spaces.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float64),
-                }
-            )
-        else:
-            raise ValueError(f"Unsupported obs_type: {self.obs_type}")
-
-        self.action_space = spaces.Box(low=ACTION_LOW, high=ACTION_HIGH, dtype=np.float32)
-
-    # Max attempts to rebuild the underlying env when MuJoCo throws
-    # `PhysicsError` (e.g. mjWARN_BADQACC) during VLABench's 20-step
-    # reset warm-up. Some random task/layout samples land in unstable
-    # initial configurations; re-sampling the layout almost always
-    # gives a stable one. A handful of upstream tasks (notably
-    # `select_mahjong`) have layout samplers that diverge often enough
-    # to need >>5 retries, so we pick a generous ceiling.
-    _ENSURE_ENV_MAX_ATTEMPTS = 20
-
-    def _ensure_env(self) -> None:
-        """Create the underlying VLABench env on first use.
-
-        Called inside the worker subprocess after fork(), so each worker gets
-        its own clean rendering context rather than inheriting a stale one from
-        the parent process (which causes crashes with AsyncVectorEnv).
-
-        Retries on `PhysicsError`: VLABench's `LM4ManipDMEnv.reset()` runs 20
-        warm-up `step()` calls while toggling gravity/fluids to let the scene
-        settle; for some random layouts MuJoCo's integrator diverges and
-        raises `mjWARN_BADQACC`. Re-sampling the layout almost always yields
-        a stable one, so we retry a number of times before giving up. Between
-        attempts we reseed NumPy's global RNG from OS entropy so the upstream
-        task sampler explores fresh initial states — without this, retries
-        can replay the same diverging configuration when the sampler is
-        deterministic given the current RNG state.
-        """
-        if self._env is not None:
-            return
-
-        import VLABench.robots  # noqa: F401  # type: ignore[import-untyped]
-        import VLABench.tasks  # noqa: F401  # type: ignore[import-untyped]
-        from dm_control.rl.control import PhysicsError  # type: ignore[import-untyped]
-        from VLABench.envs import load_env  # type: ignore[import-untyped]
-
-        h, w = self.render_resolution
-        last_exc: PhysicsError | None = None
-        for attempt in range(1, self._ENSURE_ENV_MAX_ATTEMPTS + 1):
-            try:
-                env = load_env(task=self.task, robot=self.robot, render_resolution=(h, w))
-                self._env = env
-                break
-            except PhysicsError as exc:
-                last_exc = exc
-                logger.warning(
-                    "PhysicsError on attempt %d/%d while building task '%s': %s. Retrying with fresh layout…",
-                    attempt,
-                    self._ENSURE_ENV_MAX_ATTEMPTS,
-                    self.task,
-                    exc,
-                )
-                np.random.seed(None)
-        if self._env is None:
-            assert last_exc is not None
-            raise RuntimeError(
-                f"VLABench task '{self.task}' failed to produce a stable "
-                f"initial layout after {self._ENSURE_ENV_MAX_ATTEMPTS} "
-                f"attempts. This task's upstream sampler diverges too "
-                f"often for the configured robot; consider removing it "
-                f"from the eval set. Last physics error: {last_exc}"
-            ) from last_exc
-
-        # Extract task description from the dm_control task
-        task_obj = self._env.task
-        if hasattr(task_obj, "task_description"):
-            self.task_description = task_obj.task_description
-        elif hasattr(task_obj, "language_instruction"):
-            self.task_description = task_obj.language_instruction
-        else:
-            self.task_description = self.task
-
-        # Cache robot base world position so `_build_ctrl_from_action` and
-        # `_get_obs` can translate between robot-frame (dataset) and
-        # world-frame (dm_control) without hitting physics every call.
-        try:
-            self._robot_base_xyz = np.asarray(self._env.get_robot_frame_position(), dtype=np.float64).reshape(
-                3
-            )
-        except Exception:
-            # Fallback to VLABench's default Franka base position.
-            self._robot_base_xyz = np.array([0.0, -0.4, 0.78], dtype=np.float64)
-
-    def _get_obs(self) -> dict:
-        """Get current observation from the environment."""
-        assert self._env is not None
-
-        obs = self._env.get_observation()
-        h, w = self.render_resolution
-
-        def _to_hwc3(arr: np.ndarray) -> np.ndarray:
-            """Coerce any camera array to the declared (h, w, 3) uint8 shape."""
-            a = np.asarray(arr)
-            # Drop a leading singleton batch dim if present.
-            while a.ndim > 3 and a.shape[0] == 1:
-                a = a[0]
-            if a.ndim == 3 and a.shape[0] in (1, 3, 4) and a.shape[-1] not in (1, 3, 4):
-                # CHW → HWC
-                a = np.transpose(a, (1, 2, 0))
-            if a.ndim == 2:
-                a = np.stack([a] * 3, axis=-1)
-            if a.ndim != 3:
-                return np.zeros((h, w, 3), dtype=np.uint8)
-            # Force 3 channels.
-            if a.shape[-1] == 1:
-                a = np.repeat(a, 3, axis=-1)
-            elif a.shape[-1] == 4:
-                a = a[..., :3]
-            elif a.shape[-1] != 3:
-                return np.zeros((h, w, 3), dtype=np.uint8)
-            if a.shape[:2] != (h, w):
-                a = cv2.resize(a, (w, h), interpolation=cv2.INTER_AREA)
-            return a.astype(np.uint8)
-
-        # Extract camera images — VLABench returns (n_cameras, C, H, W) or individual arrays
-        raw_frames: list[np.ndarray] = []
-        if "rgb" in obs:
-            rgb = obs["rgb"]
-            if isinstance(rgb, np.ndarray):
-                if rgb.ndim == 4:
-                    raw_frames = [rgb[i] for i in range(rgb.shape[0])]
-                elif rgb.ndim == 3:
-                    raw_frames = [rgb]
-
-        image_keys = ["image", "second_image", "wrist_image"]
-        images: dict[str, np.ndarray] = {}
-        for i, key in enumerate(image_keys):
-            if i < len(raw_frames):
-                images[key] = _to_hwc3(raw_frames[i])
-            else:
-                images[key] = np.zeros((h, w, 3), dtype=np.uint8)
-
-        # Convert VLABench's raw ee_state `[pos_world(3), quat_wxyz(4), open(1)]`
-        # to the dataset's observation.state layout `[pos_robot(3), euler_xyz(3),
-        # gripper(1)]`. See VLABench/scripts/convert_to_lerobot.py — positions
-        # are stored in robot-base frame and orientations as scipy extrinsic
-        # 'xyz' euler angles.
-        raw = np.asarray(obs.get("ee_state", np.zeros(8)), dtype=np.float64).ravel()
-        pos_world = raw[:3] if raw.size >= 3 else np.zeros(3, dtype=np.float64)
-        quat_wxyz = raw[3:7] if raw.size >= 7 else np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
-        gripper = float(raw[7]) if raw.size >= 8 else 0.0
-
-        base = self._robot_base_xyz if self._robot_base_xyz is not None else np.zeros(3, dtype=np.float64)
-        pos_robot = pos_world - base
-        euler_xyz = Rotation.from_quat([quat_wxyz[1], quat_wxyz[2], quat_wxyz[3], quat_wxyz[0]]).as_euler(
-            "xyz", degrees=False
-        )
-
-        ee_state = np.concatenate([pos_robot, euler_xyz, [gripper]]).astype(np.float64)
-
-        if self.obs_type == "pixels":
-            return {"pixels": images}
-        elif self.obs_type == "pixels_agent_pos":
-            return {
-                "pixels": images,
-                "agent_pos": ee_state.astype(np.float64),
-            }
-        else:
-            raise ValueError(f"Unknown obs_type: {self.obs_type}")
-
-    # ---- Action adaptation (EEF → joint ctrl) --------------------------------
-    #
-    # The HF vlabench datasets log 7D actions
-    # `[x, y, z (robot frame), rx, ry, rz (scipy extrinsic xyz), gripper]`,
-    # exactly matching VLABench's own eval pipeline (evaluator.base):
-    #   pos, euler, g = policy(...)
-    #   quat = euler_to_quaternion(*euler)      # extrinsic xyz -> wxyz
-    #   _, qpos = robot.get_qpos_from_ee_pos(physics, pos=pos + base, quat=quat)
-    #   env.step(np.concatenate([qpos, [g, g]]))
-    #
-    # VLABench's dm_control task writes `data.ctrl[:] = action` directly — for
-    # Franka that's 9 entries (7 arm joints + 2 gripper fingers). We mirror the
-    # above conversion so the policy's EEF commands actually drive the robot.
-
-    _FRANKA_FINGER_OPEN = 0.04  # qpos when gripper fully open
-
-    def _build_ctrl_from_action(self, action: np.ndarray, ctrl_dim: int) -> np.ndarray:
-        """Convert a 7D EEF action into the `ctrl_dim`-sized joint command vector.
-
-        For the Franka default (ctrl_dim=9): 7 arm joint qposes (via IK) +
-        2 gripper finger qposes (open/closed based on the gripper scalar).
-        If the action is already joint-space (shape matches ctrl_dim), pass
-        through.
-        """
-        if action.shape[0] == ctrl_dim:
-            return action.astype(np.float64, copy=False)
-
-        if action.shape[0] != 7:
-            # Unknown layout — fall back to zero-pad so the sim doesn't crash.
-            padded = np.zeros(ctrl_dim, dtype=np.float64)
-            padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
-            return padded
-
-        from dm_control.utils.inverse_kinematics import qpos_from_site_pose
-
-        # Action position is in robot-base frame (see convert_to_lerobot.py);
-        # dm_control's IK expects a world-frame target.
-        base = self._robot_base_xyz if self._robot_base_xyz is not None else np.zeros(3, dtype=np.float64)
-        pos_world = np.asarray(action[:3], dtype=np.float64) + base
-        rx, ry, rz = float(action[3]), float(action[4]), float(action[5])
-        gripper = float(np.clip(action[6], 0.0, 1.0))
-
-        # Dataset euler is scipy extrinsic 'xyz' (same as VLABench's
-        # `euler_to_quaternion`). scipy emits `[x, y, z, w]`; dm_control's IK
-        # and MuJoCo use `[w, x, y, z]`, so reorder.
-        qxyzw = Rotation.from_euler("xyz", [rx, ry, rz], degrees=False).as_quat()
-        quat = np.array([qxyzw[3], qxyzw[0], qxyzw[1], qxyzw[2]], dtype=np.float64)
-
-        assert self._env is not None
-        robot = self._env.task.robot
-        site_name = robot.end_effector_site.full_identifier
-
-        # inplace=False so IK doesn't mutate physics state mid-step — we only
-        # want the solved qpos. Fetch a fresh physics handle — caching it can
-        # yield a stale weakref after a reset.
-        ik_result = qpos_from_site_pose(
-            self._env.physics,
-            site_name=site_name,
-            target_pos=pos_world,
-            target_quat=quat,
-            inplace=False,
-            max_steps=100,
-        )
-        n_dof = robot.n_dof  # 7 for Franka
-        arm_qpos = ik_result.qpos[:n_dof]
-
-        # Dataset gripper convention: 1 = open (finger qpos = 0.04),
-        # 0 = closed (finger qpos = 0.0). See VLABench/scripts/convert_to_lerobot.py
-        # where `trajectory[i][-1] > 0.03` is encoded as `1`.
-        finger_qpos = gripper * self._FRANKA_FINGER_OPEN
-
-        ctrl = np.zeros(ctrl_dim, dtype=np.float64)
-        ctrl[:n_dof] = arm_qpos
-        # Remaining entries are gripper fingers (usually 2 for Franka).
-        ctrl[n_dof:] = finger_qpos
-        return ctrl
-
-    def reset(self, seed=None, **kwargs) -> tuple[RobotObservation, dict[str, Any]]:
-        self._ensure_env()
-        assert self._env is not None
-        super().reset(seed=seed)
-
-        if seed is not None:
-            self._seed_inner_env(int(self.np_random.integers(0, 2**31 - 1)))
-
-        self._env.reset()
-
-        observation = self._get_obs()
-        info = {"is_success": False}
-        return observation, info
-
-    def _seed_inner_env(self, seed: int) -> None:
-        """Propagate `seed` to the inner dm_control env. `Environment.reset()`
-        doesn't accept a seed, so we re-seed the task and environment
-        `RandomState`s directly. Best-effort: silently skipped when the
-        expected attributes are absent on a given VLABench version.
-        """
-        for owner_attr, rng_attr in (("task", "random"), (None, "_random_state")):
-            owner = getattr(self._env, owner_attr) if owner_attr else self._env
-            rng = getattr(owner, rng_attr, None)
-            rng_seed = getattr(rng, "seed", None)
-            if callable(rng_seed):
-                rng_seed(seed)
-
-    def step(self, action: np.ndarray) -> tuple[RobotObservation, float, bool, bool, dict[str, Any]]:
-        from dm_control.rl.control import PhysicsError  # type: ignore[import-untyped]
-
-        self._ensure_env()
-        assert self._env is not None
-
-        if action.ndim != 1:
-            raise ValueError(
-                f"Expected action to be 1-D (shape (action_dim,)), "
-                f"but got shape {action.shape} with ndim={action.ndim}"
-            )
-
-        if self.action_mode not in ("eef", "joint", "delta_eef"):
-            raise ValueError(f"Unknown action_mode: {self.action_mode}")
-
-        # Always refetch physics — dm_control returns a weakref proxy that can
-        # go stale across resets.
-        physics = self._env.physics
-        ctrl_dim = int(physics.data.ctrl.shape[0])
-        ctrl = self._build_ctrl_from_action(action, ctrl_dim)
-        try:
-            timestep = self._env.step(ctrl)
-        except PhysicsError as exc:
-            # Physics integrator diverged (e.g. mjWARN_BADQACC). Treat it as
-            # a graceful failed termination rather than a hard crash — the
-            # rest of the multi-task eval should still run.
-            logger.warning(
-                "PhysicsError during step on task '%s': %s. Terminating episode.",
-                self.task,
-                exc,
-            )
-            observation = self._get_obs()
-            info = {"task": self.task, "is_success": False, "physics_error": True}
-            # Drop the stale env so the next reset() rebuilds it cleanly.
-            with contextlib.suppress(Exception):
-                self._env.close()
-            self._env = None
-            return observation, 0.0, True, False, info
-
-        # Extract reward from dm_control timestep
-        reward = float(timestep.reward) if timestep.reward is not None else 0.0
-
-        # Check success via the task's termination condition
-        is_success = False
-        if hasattr(self._env, "task") and hasattr(self._env.task, "should_terminate_episode"):
-            is_success = bool(self._env.task.should_terminate_episode(self._env.physics))
-
-        terminated = is_success
-        truncated = False
-        info = {
-            "task": self.task,
-            "is_success": is_success,
-        }
-
-        observation = self._get_obs()
-
-        if terminated:
-            self.reset()
-
-        return observation, reward, terminated, truncated, info
-
-    def render(self) -> np.ndarray:
-        self._ensure_env()
-        obs = self._get_obs()
-        return obs["pixels"]["image"]
-
-    def close(self):
-        if self._env is not None:
-            self._env.close()
-            self._env = None
-
-
-# ---- Main API ----------------------------------------------------------------
-
-
-def create_vlabench_envs(
-    task: str,
-    n_envs: int,
-    gym_kwargs: dict[str, Any] | None = None,
-    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
-) -> dict[str, dict[int, Any]]:
-    """
-    Create vectorized VLABench environments with a consistent return shape.
-
-    Returns:
-        dict[suite_name][task_id] -> vec_env (env_cls([...]) with exactly n_envs factories)
-
-    Notes:
-        - n_envs is the number of rollouts *per task*.
-        - `task` can be a suite name ("primitive", "composite"), a comma-separated list of
-          suite names, or individual task names (e.g. "select_fruit,heat_food").
-    """
-    if env_cls is None or not callable(env_cls):
-        raise ValueError("env_cls must be a callable that wraps a list of environment factory callables.")
-    if not isinstance(n_envs, int) or n_envs <= 0:
-        raise ValueError(f"n_envs must be a positive int; got {n_envs}.")
-
-    gym_kwargs = dict(gym_kwargs or {})
-    task_groups = [t.strip() for t in task.split(",") if t.strip()]
-    if not task_groups:
-        raise ValueError("`task` must contain at least one VLABench task or suite name.")
-
-    logger.info(
-        "Creating VLABench envs | task_groups=%s | n_envs(per task)=%d",
-        task_groups,
-        n_envs,
-    )
-
-    is_async = env_cls is gym.vector.AsyncVectorEnv
-    cached_obs_space = None
-    cached_act_space = None
-    cached_metadata = None
-    out: dict[str, dict[int, Any]] = defaultdict(dict)
-
-    for group in task_groups:
-        # Check if it's a suite name, otherwise treat as individual task
-        tasks = SUITE_TASKS.get(group, [group])
-
-        for tid, task_name in enumerate(tasks):
-            logger.info(
-                "Building vec env | group=%s | task_id=%d | task=%s",
-                group,
-                tid,
-                task_name,
-            )
-
-            fns = [(lambda tn=task_name: VLABenchEnv(task=tn, **gym_kwargs)) for _ in range(n_envs)]
-
-            if is_async:
-                lazy = _LazyAsyncVectorEnv(fns, cached_obs_space, cached_act_space, cached_metadata)
-                if cached_obs_space is None:
-                    cached_obs_space = lazy.observation_space
-                    cached_act_space = lazy.action_space
-                    cached_metadata = lazy.metadata
-                out[group][tid] = lazy
-            else:
-                out[group][tid] = env_cls(fns)
-
-    return {group: dict(task_map) for group, task_map in out.items()}

src/lerobot/policies/__init__.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from lerobot.utils.action_interpolator import ActionInterpolator as ActionInterpolator
+
 from .act.configuration_act import ACTConfig as ACTConfig
 from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
 from .factory import get_policy_class, make_policy, make_policy_config, make_pre_post_processors
@@ -21,7 +23,6 @@ from .pi0.configuration_pi0 import PI0Config as PI0Config
 from .pi0_fast.configuration_pi0_fast import PI0FastConfig as PI0FastConfig
 from .pi05.configuration_pi05 import PI05Config as PI05Config
 from .pretrained import PreTrainedPolicy as PreTrainedPolicy
-from .rtc import ActionInterpolator as ActionInterpolator
 from .sac.configuration_sac import SACConfig as SACConfig
 from .sac.reward_model.configuration_classifier import RewardClassifierConfig as RewardClassifierConfig
 from .sarm.configuration_sarm import SARMConfig as SARMConfig

src/lerobot/policies/factory.py

@@ -140,10 +140,6 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from .smolvla.modeling_smolvla import SmolVLAPolicy
 
         return SmolVLAPolicy
-    elif name == "smolvla2":
-        from .smolvla2.modeling_smolvla2 import SmolVLA2Policy
-
-        return SmolVLA2Policy
     elif name == "sarm":
         from .sarm.modeling_sarm import SARMRewardModel
 
@@ -204,10 +200,6 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return SACConfig(**kwargs)
     elif policy_type == "smolvla":
         return SmolVLAConfig(**kwargs)
-    elif policy_type == "smolvla2":
-        from .smolvla2.configuration_smolvla2 import SmolVLA2Config
-
-        return SmolVLA2Config(**kwargs)
     elif policy_type == "reward_classifier":
         return RewardClassifierConfig(**kwargs)
     elif policy_type == "groot":
@@ -394,17 +386,6 @@ def make_pre_post_processors(
             dataset_stats=kwargs.get("dataset_stats"),
         )
 
-    elif policy_cfg.type == "smolvla2":
-        # NOTE: SmolVLA2Config subclasses SmolVLAConfig, so this branch
-        # MUST come before the SmolVLAConfig isinstance check below
-        # (otherwise SmolVLA2 would silently pick up SmolVLA's processor).
-        from .smolvla2.processor_smolvla2 import make_smolvla2_pre_post_processors
-
-        processors = make_smolvla2_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
     elif isinstance(policy_cfg, SmolVLAConfig):
         from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
 

src/lerobot/policies/rtc/action_interpolator.py

@@ -1,116 +1,4 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
+# Moved to lerobot.utils.action_interpolator — re-exported for backwards compatibility.
+from lerobot.utils.action_interpolator import ActionInterpolator
 
-"""Action interpolation for smoother robot control.
-
-Provides configurable Nx control rate by interpolating between consecutive actions.
-Useful with RTC and action-chunking policies to reduce jerkiness.
-"""
-
-from torch import Tensor
-
-
-class ActionInterpolator:
-    """Interpolates between consecutive actions for smoother control.
-
-    When enabled with multiplier N, produces N actions per policy action
-    by linearly interpolating between the previous and current action.
-
-    Example with multiplier=3:
-        prev_action -> [1/3 interpolated, 2/3 interpolated, current_action]
-
-    This effectively multiplies the control rate for smoother motion.
-
-    Usage:
-        interpolator = ActionInterpolator(multiplier=2)  # 2x control rate
-
-        # In control loop:
-        if interpolator.needs_new_action():
-            new_action = queue.get()
-            if new_action:
-                interpolator.add(new_action.cpu())
-
-        action = interpolator.get()
-        if action:
-            robot.send_action(action)
-    """
-
-    def __init__(self, multiplier: int = 1):
-        """Initialize the interpolator.
-
-        Args:
-            multiplier: Control rate multiplier (1 = no interpolation, 2 = 2x, 3 = 3x, etc.)
-        """
-        if multiplier < 1:
-            raise ValueError(f"multiplier must be >= 1, got {multiplier}")
-        self.multiplier = multiplier
-        self._prev: Tensor | None = None
-        self._buffer: list[Tensor] = []
-        self._idx = 0
-
-    @property
-    def enabled(self) -> bool:
-        """Whether interpolation is active (multiplier > 1)."""
-        return self.multiplier > 1
-
-    def reset(self):
-        """Reset interpolation state (call between episodes)."""
-        self._prev = None
-        self._buffer = []
-        self._idx = 0
-
-    def needs_new_action(self) -> bool:
-        """Check if a new action is needed from the queue."""
-        return self._idx >= len(self._buffer)
-
-    def add(self, action: Tensor) -> None:
-        """Add a new action and compute interpolated sequence.
-
-        Args:
-            action: New action tensor from policy/queue (already on CPU).
-        """
-        if self.multiplier > 1 and self._prev is not None:
-            self._buffer = []
-            for i in range(1, self.multiplier + 1):
-                t = i / self.multiplier
-                interp = self._prev + t * (action - self._prev)
-                self._buffer.append(interp)
-        else:
-            # First step: no previous action yet, so run at base FPS without interpolation.
-            self._buffer = [action.clone()]
-        self._prev = action.clone()
-        self._idx = 0
-
-    def get(self) -> Tensor | None:
-        """Get the next interpolated action.
-
-        Returns:
-            Next action tensor, or None if buffer is exhausted.
-        """
-        if self._idx >= len(self._buffer):
-            return None
-        action = self._buffer[self._idx]
-        self._idx += 1
-        return action
-
-    def get_control_interval(self, fps: float) -> float:
-        """Get the control interval based on interpolation multiplier.
-
-        Args:
-            fps: Base frames per second.
-
-        Returns:
-            Control interval in seconds (divided by multiplier).
-        """
-        return 1.0 / (fps * self.multiplier)
+__all__ = ["ActionInterpolator"]

src/lerobot/policies/rtc/action_queue.py

@@ -92,10 +92,10 @@ class ActionQueue:
         Returns:
             int: Number of unconsumed actions.
         """
-        if self.queue is None:
-            return 0
-        length = len(self.queue)
-        return length - self.last_index
+        with self.lock:
+            if self.queue is None:
+                return 0
+            return len(self.queue) - self.last_index
 
     def empty(self) -> bool:
         """Check if the queue is empty.
@@ -103,11 +103,10 @@ class ActionQueue:
         Returns:
             bool: True if no actions remain, False otherwise.
         """
-        if self.queue is None:
-            return True
-
-        length = len(self.queue)
-        return length - self.last_index <= 0
+        with self.lock:
+            if self.queue is None:
+                return True
+            return len(self.queue) - self.last_index <= 0
 
     def get_action_index(self) -> int:
         """Get the current action consumption index.
@@ -115,7 +114,8 @@ class ActionQueue:
         Returns:
             int: Index of the next action to be consumed.
         """
-        return self.last_index
+        with self.lock:
+            return self.last_index
 
     def get_left_over(self) -> Tensor | None:
         """Get leftover original actions for RTC prev_chunk_left_over.

src/lerobot/policies/smolvla2/__init__.py

@@ -1,38 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""SmolVLA2 — SmolVLA with the SmolVLM language head re-enabled.
-
-SmolVLA strips the LM head from the SmolVLM backbone because it only does
-flow-matching action prediction. SmolVLA2 keeps the LM head so the same
-model can train on the full Hi Robot / MEM / ECoT message blend defined in
-the steerable annotation plan (PR1 + PR2):
-
-* action-only sub-recipes (e.g. ``low_level_execution``) → flow loss
-* text-only sub-recipes (e.g. ``memory_update``, ``ask_vqa``,
-  ``user_interjection_response``, ``high_level_subtask``) → CE loss on
-  ``lm_head`` over the recipe's target message tokens
-* mixed sub-recipes → both losses summed (weighted)
-
-The ``predict_actions`` toggle follows the Pi0.5 convention from Section
-I.7 of the plan: ``True`` if any ``low_level`` target is present in the
-sample, else ``False``.
-
-This package is a thin subclass of ``lerobot.policies.smolvla`` so most of
-the model code stays in one place — only the dual-loss path and the
-chat-template processor live here.
-"""
-
-from .configuration_smolvla2 import SmolVLA2Config
-
-__all__ = ["SmolVLA2Config"]

src/lerobot/policies/smolvla2/chat_processor_smolvla2.py

@@ -1,271 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""SmolVLA2's chat-template tokenization step.
-
-Replaces SmolVLA's plain ``TokenizerProcessorStep`` for SmolVLA2 when a
-``recipe_path`` is set. Reads the rendered messages produced by
-``RenderMessagesStep`` (PR 1) and produces:
-
-* ``OBS_LANGUAGE_TOKENS`` / ``OBS_LANGUAGE_ATTENTION_MASK`` —
-  the chat-templated prompt tokenized by SmolVLM's tokenizer, with
-  ``tools=meta.tools`` (PR 1's catalog).
-* ``text_labels`` — same shape as token ids, ``-100`` everywhere except
-  the positions belonging to messages whose index is in
-  ``target_message_indices``. The next commit's modeling forward path
-  applies cross-entropy on those positions via the SmolVLM ``lm_head``.
-* ``predict_actions`` — bool tensor, ``True`` iff any of the rendered
-  target messages has ``message_streams[i] == "low_level"``. The
-  modeling forward uses this to gate the flow head.
-
-Image / video content blocks in the rendered messages are dropped
-before tokenization — the chat template only handles text, and SmolVLA
-already passes camera tensors out-of-band via the standard
-``OBS_IMAGES_*`` features. This keeps the prefix layout unchanged
-(``embed_prefix`` puts image embeddings before language embeddings,
-matching the chat-template-stripped text order).
-"""
-
-from __future__ import annotations
-
-import copy
-import logging
-from dataclasses import dataclass
-from typing import Any
-
-import torch
-
-from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.datasets.language import DEFAULT_TOOLS
-from lerobot.processor.pipeline import ProcessorStep, ProcessorStepRegistry
-from lerobot.types import EnvTransition, TransitionKey
-from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="smolvla2_chat_tokenizer")
-class SmolVLA2ChatTokenizerStep(ProcessorStep):
-    """Render messages → token ids + label mask + predict_actions flag.
-
-    This is the bridge between the recipe stack (PR 1's
-    ``RenderMessagesStep`` outputs) and the SmolVLA2 modeling forward
-    (next commit, which reads ``text_labels`` / ``predict_actions``).
-    Pure-text turns and multi-stream targets are both handled.
-    """
-
-    tokenizer_name: str = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
-    max_length: int = 2048
-    padding: str = "longest"
-    padding_side: str = "right"
-    tools: list[dict[str, Any]] | None = None
-
-    def __post_init__(self) -> None:
-        # Lazy: don't load the tokenizer until the step actually runs,
-        # so unit tests that import the module without transformers
-        # installed still pass.
-        self._tokenizer: Any = None
-        if self.tools is None:
-            # Default: ship the canonical ``say`` schema. Users who set
-            # ``meta.tools`` differently can override via
-            # ``with_tools(meta.tools)``.
-            self.tools = list(DEFAULT_TOOLS)
-
-    # ------------------------------------------------------------------
-    # Public API
-    # ------------------------------------------------------------------
-
-    def with_tools(self, tools: list[dict[str, Any]]) -> "SmolVLA2ChatTokenizerStep":
-        """Override the tools catalog rendered into the system prompt."""
-        self.tools = list(tools)
-        return self
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition | None:
-        comp = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
-        messages = comp.get("messages")
-        if not messages:
-            # No recipe rendering happened — nothing to do; downstream
-            # falls back to whatever ``task`` is in the transition.
-            return transition
-
-        message_streams: list[str | None] = list(comp.get("message_streams") or [])
-        target_indices: list[int] = sorted(
-            int(i) for i in (comp.get("target_message_indices") or [])
-        )
-
-        tokenizer = self._get_tokenizer()
-        text_messages = [_strip_lerobot_blocks(m) for m in messages]
-
-        # Tokenize the full chat once.
-        full_ids = tokenizer.apply_chat_template(
-            text_messages,
-            tools=self.tools,
-            add_generation_prompt=False,
-            tokenize=True,
-            return_tensors=None,
-        )
-        if isinstance(full_ids, list) and full_ids and isinstance(full_ids[0], list):
-            full_ids = full_ids[0]
-
-        # Build the label mask by re-rendering progressively up to each
-        # target message and reading off the prefix length. This is the
-        # robust way to get exact token boundaries: we use the same
-        # tokenizer, the same ``tools=`` argument, and the same chat
-        # template — so the prefix tokens are guaranteed to be a prefix
-        # of the full sequence.
-        labels = [-100] * len(full_ids)
-        for tgt in target_indices:
-            prefix_ids = tokenizer.apply_chat_template(
-                text_messages[:tgt],
-                tools=self.tools,
-                add_generation_prompt=False,
-                tokenize=True,
-                return_tensors=None,
-            )
-            full_through_target = tokenizer.apply_chat_template(
-                text_messages[: tgt + 1],
-                tools=self.tools,
-                add_generation_prompt=False,
-                tokenize=True,
-                return_tensors=None,
-            )
-            if isinstance(prefix_ids, list) and prefix_ids and isinstance(prefix_ids[0], list):
-                prefix_ids = prefix_ids[0]
-            if (
-                isinstance(full_through_target, list)
-                and full_through_target
-                and isinstance(full_through_target[0], list)
-            ):
-                full_through_target = full_through_target[0]
-            start = len(prefix_ids)
-            end = min(len(full_through_target), len(full_ids))
-            for pos in range(start, end):
-                labels[pos] = int(full_ids[pos])
-
-        # Truncate / pad to ``max_length`` so batches collate cleanly.
-        # The SmolVLA pipeline downstream relies on a fixed length
-        # behaviour ("longest" or "max_length") — we mirror it here.
-        if len(full_ids) > self.max_length:
-            full_ids = full_ids[: self.max_length]
-            labels = labels[: self.max_length]
-        attn = [1] * len(full_ids)
-        if self.padding == "max_length" and len(full_ids) < self.max_length:
-            pad_id = (
-                tokenizer.pad_token_id
-                if tokenizer.pad_token_id is not None
-                else 0
-            )
-            n_pad = self.max_length - len(full_ids)
-            full_ids = full_ids + [pad_id] * n_pad
-            labels = labels + [-100] * n_pad
-            attn = attn + [0] * n_pad
-
-        ids_t = torch.tensor(full_ids, dtype=torch.long)
-        attn_t = torch.tensor(attn, dtype=torch.bool)
-        labels_t = torch.tensor(labels, dtype=torch.long)
-        predict_actions = any(
-            i < len(message_streams) and message_streams[i] == "low_level"
-            for i in target_indices
-        )
-
-        new_complementary = dict(comp)
-        # Drop the per-recipe sidecar keys; everything downstream needs
-        # is now in the tokenized form.
-        new_complementary.pop("messages", None)
-        new_complementary.pop("message_streams", None)
-        new_complementary.pop("target_message_indices", None)
-        # SmolVLA's pipeline expects ``OBS_LANGUAGE_TOKENS`` /
-        # ``OBS_LANGUAGE_ATTENTION_MASK`` on the OBSERVATION key. Place
-        # them there — and drop ``task`` so the upstream
-        # ``TokenizerProcessorStep`` (which we replace) doesn't double-
-        # tokenize.
-        observation = dict(transition.get(TransitionKey.OBSERVATION) or {})
-        observation[OBS_LANGUAGE_TOKENS] = ids_t
-        observation[OBS_LANGUAGE_ATTENTION_MASK] = attn_t
-        new_complementary["text_labels"] = labels_t
-        new_complementary["predict_actions"] = torch.tensor(predict_actions, dtype=torch.bool)
-        new_complementary.pop("task", None)
-
-        new_transition = dict(transition)
-        new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary
-        new_transition[TransitionKey.OBSERVATION] = observation
-        return new_transition
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        """Pass-through; this step writes runtime tensors not features."""
-        return features
-
-    # ------------------------------------------------------------------
-    # Helpers
-    # ------------------------------------------------------------------
-
-    def _get_tokenizer(self):  # noqa: ANN202
-        if self._tokenizer is not None:
-            return self._tokenizer
-        try:
-            from transformers import AutoTokenizer  # noqa: PLC0415
-        except ImportError as exc:  # pragma: no cover
-            raise ImportError(
-                "SmolVLA2ChatTokenizerStep requires transformers. "
-                "`pip install lerobot[transformers-dep]`."
-            ) from exc
-        self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
-        if self._tokenizer.pad_token_id is None and self._tokenizer.eos_token_id is not None:
-            self._tokenizer.pad_token = self._tokenizer.eos_token
-        return self._tokenizer
-
-
-def _strip_lerobot_blocks(message: dict[str, Any]) -> dict[str, Any]:
-    """Remove LeRobot-specific multimodal blocks from ``message`` content.
-
-    The recipe DSL allows authors to write multimodal content like
-    ``{"type": "image", "feature": "observation.images.top"}``. SmolVLM's
-    tokenizer doesn't know that ``feature`` key (it expects ``url`` or
-    ``path``). The actual image tensor flows through SmolVLA's
-    ``OBS_IMAGES_*`` channels separately; the chat template only needs
-    the text. So we strip non-text blocks before tokenizing.
-    """
-    new = dict(message)
-    content = new.get("content")
-    if isinstance(content, list):
-        text_parts: list[dict[str, Any]] = []
-        for block in content:
-            if not isinstance(block, dict):
-                continue
-            if block.get("type") == "text":
-                text_parts.append({"type": "text", "text": str(block.get("text", ""))})
-        # If only one text block survives, flatten to a string for
-        # template friendliness; some chat templates choke on a single-
-        # element list.
-        if len(text_parts) == 1:
-            new["content"] = text_parts[0]["text"]
-        elif text_parts:
-            new["content"] = text_parts
-        else:
-            new["content"] = ""
-    if "tool_calls" in new and not new["tool_calls"]:
-        # Drop empty tool_calls — some templates render them as a
-        # spurious empty marker.
-        new.pop("tool_calls")
-    # ``stream`` and ``target`` were recipe metadata; templates don't
-    # know them and may warn or crash.
-    new.pop("stream", None)
-    new.pop("target", None)
-    return new
-
-
-# Re-export for tests / introspection
-strip_lerobot_blocks = _strip_lerobot_blocks

src/lerobot/policies/smolvla2/configuration_smolvla2.py

@@ -1,97 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-
-from lerobot.configs import PreTrainedConfig
-
-from ..smolvla.configuration_smolvla import SmolVLAConfig
-
-
-@PreTrainedConfig.register_subclass("smolvla2")
-@dataclass
-class SmolVLA2Config(SmolVLAConfig):
-    """SmolVLA2 — SmolVLA with the underlying SmolVLM language head re-enabled.
-
-    SmolVLA strips the LM head from the SmolVLM backbone because it only
-    needs flow-matching action prediction. SmolVLA2 keeps the LM head so the
-    same model can train on:
-
-      * **action-only sub-recipes** (e.g. ``low_level_execution``) — flow loss
-        on the action expert, same as SmolVLA. ``predict_actions=True``.
-      * **text-only sub-recipes** (e.g. ``memory_update`` / ``ask_vqa`` /
-        ``user_interjection_response`` / ``high_level_subtask``) — cross-
-        entropy loss on the LM head over the recipe's target message tokens.
-        Skips the flow head entirely. ``predict_actions=False``.
-      * **mixed sub-recipes** — both heads run, losses summed (weighted).
-
-    The split is controlled by ``predict_actions = bool(targets_by_stream
-    .get("low_level"))`` per the Pi0.5 convention in the steerable
-    annotation plan (Section I.7), implemented inside the processor /
-    forward path. Recipes drive it via ``stream`` + ``target`` metadata.
-
-    Compared to ``SmolVLAConfig`` this adds:
-
-    - ``recipe_path``: path to a ``TrainingRecipe`` YAML (loaded by the
-      train script). When ``None``, SmolVLA2 falls back to the SmolVLA
-      task-only path so unannotated datasets still work.
-    - ``text_loss_weight`` / ``flow_loss_weight``: relative weights when
-      both losses are active in a single sample.
-    - ``unfreeze_lm_head``: must be ``True`` for the text head to learn —
-      SmolVLA freezes ``lm_head`` to "avoid unused params issues" and we
-      need to undo that for SmolVLA2.
-    - ``train_expert_only=False`` by default, since the VLM body now also
-      participates in text-target gradients.
-    """
-
-    # Recipe / language stack ---------------------------------------------
-    recipe_path: str | None = "recipes/smolvla2_hirobot.yaml"
-    """Path (absolute or relative to ``src/lerobot/configs/``) to a
-    ``TrainingRecipe`` YAML. The default points at the canonical Hi Robot
-    blend shipped alongside SmolVLA2. Set to ``None`` to disable recipe
-    rendering and fall back to SmolVLA's single-task prompt path
-    (unannotated datasets keep working that way)."""
-
-    apply_chat_template: bool = True
-    """Apply the SmolVLM tokenizer's chat template to the rendered messages
-    before tokenizing. SmolVLM's backbone is chat-pretrained, so this
-    matches its training distribution."""
-
-    # Loss weights --------------------------------------------------------
-    text_loss_weight: float = 1.0
-    """Weight on the LM-head cross-entropy term. Set to ``0`` to disable
-    text training entirely (reverts to flow-only / SmolVLA behaviour)."""
-
-    flow_loss_weight: float = 1.0
-    """Weight on the action-expert flow-matching term."""
-
-    # Backbone training ---------------------------------------------------
-    unfreeze_lm_head: bool = True
-    """Whether to unfreeze the SmolVLM ``lm_head`` (and the immediately
-    preceding norm + last text-model layer that SmolVLA freezes). Must be
-    ``True`` for the text head to learn. Setting this to ``False``
-    effectively reduces SmolVLA2 back to SmolVLA's flow-only training,
-    which is occasionally useful for ablations."""
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        # Backbone needs gradients flowing through its text path when the
-        # LM head is producing supervised text. Override the SmolVLA
-        # default (`train_expert_only=True`) unless the user explicitly
-        # opts out of text training via `text_loss_weight=0`.
-        if self.text_loss_weight > 0 and self.unfreeze_lm_head:
-            # The user can still flip this back via CLI; this only
-            # changes the *default* when SmolVLA2 is actually training a
-            # text head.
-            self.train_expert_only = False

src/lerobot/policies/smolvla2/inference/__init__.py

@@ -1,68 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""SmolVLA2 inference / runtime orchestration.
-
-Multi-rate runtime that mirrors the recipe-time training shape:
-
-  low_level_execution        → LowLevelForward + DispatchAction (high Hz)
-  high_level_subtask         → HighLevelSubtaskFwd (~1 Hz)
-  memory_update              → MemoryUpdateFwd (event: subtask_change)
-  user_interjection_response → UserInterjectionFwd (event: stdin)
-  ask_vqa_*                  → AskVQAFwd (event: stdin question)
-  speech tool calls          → DispatchToolCalls (event: tool_call_pending)
-
-The CLI ``lerobot-smolvla2-runtime`` builds an ``SmolVLA2Runtime`` and
-calls ``run()``.
-"""
-
-from .repl import StdinReader
-from .runtime import SmolVLA2Runtime
-from .runtime_state import initial_runtime_state, push_log, set_if_changed, take_event
-from .steps import (
-    AskVQAFwd,
-    DispatchAction,
-    DispatchToolCalls,
-    HighLevelSubtaskFwd,
-    InferenceStep,
-    LowLevelForward,
-    MemoryUpdateFwd,
-    UserInterjectionFwd,
-)
-from .triggers import EventTrigger, HzTrigger, Tick, TickClock, Trigger
-
-__all__ = [
-    # runtime
-    "SmolVLA2Runtime",
-    "StdinReader",
-    # state helpers
-    "initial_runtime_state",
-    "push_log",
-    "set_if_changed",
-    "take_event",
-    # triggers
-    "Trigger",
-    "Tick",
-    "TickClock",
-    "HzTrigger",
-    "EventTrigger",
-    # steps
-    "InferenceStep",
-    "LowLevelForward",
-    "DispatchAction",
-    "HighLevelSubtaskFwd",
-    "MemoryUpdateFwd",
-    "UserInterjectionFwd",
-    "AskVQAFwd",
-    "DispatchToolCalls",
-]

src/lerobot/policies/smolvla2/inference/repl.py

@@ -1,87 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Stdin REPL event collector for the SmolVLA2 runtime.
-
-Reads non-blocking stdin lines, classifies each one heuristically:
-
-  "stop" / "quit" / "exit"               → state["stop"] = True
-  ends with "?"                           → user_vqa_query event
-  starts with "task:" or first line       → set runtime task
-  anything else                           → user_interjection event
-
-Plugged into the runtime via ``event_collector=StdinReader().poll``.
-"""
-
-from __future__ import annotations
-
-import select
-import sys
-from dataclasses import dataclass, field
-from typing import Any
-
-
-@dataclass
-class StdinReader:
-    """Non-blocking stdin line collector for the runtime loop."""
-
-    prompt: str = "> "
-    _seen_first_line: bool = field(default=False, init=False)
-    _prompted: bool = field(default=False, init=False)
-
-    def poll(self, state: dict[str, Any]) -> None:
-        """Drain pending stdin lines into runtime events."""
-        # Print the input prompt once on every fresh tick if we don't
-        # already have a pending line; matches the expected REPL feel.
-        if not self._prompted:
-            print(self.prompt, end="", flush=True)
-            self._prompted = True
-
-        # ``select`` with timeout=0 makes this non-blocking. Only works
-        # for actual TTY / pipe stdins; CI / scripted runs hit EOF.
-        try:
-            ready, _, _ = select.select([sys.stdin], [], [], 0)
-        except (ValueError, OSError):
-            return
-        if not ready:
-            return
-
-        line = sys.stdin.readline()
-        if not line:  # EOF
-            state["stop"] = True
-            return
-        line = line.strip()
-        self._prompted = False  # we'll re-prompt next tick
-        if not line:
-            return
-
-        lower = line.lower()
-        if lower in {"stop", "quit", "exit"}:
-            state["stop"] = True
-            return
-
-        # First non-control line sets the task if no task is active.
-        if not state.get("task"):
-            task = line[5:].strip() if lower.startswith("task:") else line
-            state["task"] = task
-            print(f"[smolvla2] Task: {task}", flush=True)
-            self._seen_first_line = True
-            return
-
-        # Question → VQA; statement → interjection.
-        if lower.endswith("?"):
-            state["recent_vqa_query"] = line
-            state.setdefault("events_this_tick", []).append("user_vqa_query")
-        else:
-            state["recent_interjection"] = line
-            state.setdefault("events_this_tick", []).append("user_interjection")

src/lerobot/policies/smolvla2/inference/runtime.py

@@ -1,143 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""SmolVLA2 runtime loop.
-
-Threads the multi-rate inference pipeline together with a stdin REPL
-event collector, drives ticks through :class:`TickClock`, and prints
-state-change updates to the user.
-"""
-
-from __future__ import annotations
-
-import logging
-from collections import deque
-from dataclasses import dataclass, field
-from typing import Any, Callable
-
-from .runtime_state import initial_runtime_state, push_log
-from .steps import (
-    AskVQAFwd,
-    DispatchAction,
-    DispatchToolCalls,
-    HighLevelSubtaskFwd,
-    InferenceStep,
-    LowLevelForward,
-    MemoryUpdateFwd,
-    UserInterjectionFwd,
-)
-from .triggers import HzTrigger, TickClock
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class SmolVLA2Runtime:
-    """Compose the inference pipeline and drive it tick-by-tick."""
-
-    policy: Any
-    tools: dict[str, Any] = field(default_factory=dict)
-    """Name → tool-instance dict, e.g. ``{"say": SayTool(...)}``. Read
-    from :func:`lerobot.tools.get_tools(meta)` when wiring the
-    runtime."""
-    observation_provider: Callable[[], dict | None] | None = None
-    """Closure returning the current preprocessed observation batch.
-    ``None`` for dry-run / language-only sessions."""
-    robot_executor: Callable[[Any], None] | None = None
-    """Closure that takes one action chunk and forwards it to the
-    robot. ``None`` for dry-run."""
-    event_collector: Callable[[dict], None] | None = None
-    """Per-tick hook that polls external sources (stdin, network) and
-    appends event names to ``state["events_this_tick"]``."""
-    chunk_hz: float = 4.0
-    ctrl_hz: float = 50.0
-    high_level_hz: float = 1.0
-    max_rate_hz: float = 50.0
-
-    pipeline: list[InferenceStep] = field(init=False)
-    state: dict[str, Any] = field(init=False)
-    _stop: bool = field(default=False, init=False)
-
-    def __post_init__(self) -> None:
-        self.pipeline = [
-            LowLevelForward(
-                trigger=HzTrigger(self.chunk_hz),
-                policy=self.policy,
-                observation_provider=self.observation_provider,
-            ),
-            DispatchAction(
-                trigger=HzTrigger(self.ctrl_hz),
-                robot_executor=self.robot_executor,
-            ),
-            HighLevelSubtaskFwd(
-                trigger=HzTrigger(self.high_level_hz),
-                policy=self.policy,
-            ),
-            MemoryUpdateFwd(policy=self.policy),
-            UserInterjectionFwd(policy=self.policy),
-            AskVQAFwd(policy=self.policy),
-            DispatchToolCalls(tools=self.tools),
-        ]
-        self.state = initial_runtime_state()
-
-    # ------------------------------------------------------------------
-    # Lifecycle
-    # ------------------------------------------------------------------
-
-    def set_task(self, task: str) -> None:
-        """Set or replace the active task. Logged for the REPL."""
-        self.state["task"] = task
-        push_log(self.state, f"Task: {task}")
-
-    def stop(self) -> None:
-        self._stop = True
-
-    def run(self, *, max_ticks: int | None = None) -> None:
-        """Main loop. Returns when ``stop()`` is called or after
-        ``max_ticks`` ticks (useful for tests / dry-run)."""
-        clock = TickClock(max_rate_hz=self.max_rate_hz)
-        while not self._stop:
-            tick = clock.advance()
-            self.state["_tick"] = tick
-            self.state["events_this_tick"] = []
-            self.state["log_lines"] = []
-
-            if self.event_collector is not None:
-                self.event_collector(self.state)
-            if self.state.get("stop"):
-                self._stop = True
-                break
-
-            for step in self.pipeline:
-                self.state = step(self.state)
-
-            self._flush_logs()
-            if max_ticks is not None and tick.index >= max_ticks:
-                break
-
-        self._on_shutdown()
-
-    # ------------------------------------------------------------------
-    # I/O
-    # ------------------------------------------------------------------
-
-    def _flush_logs(self) -> None:
-        for line in self.state.get("log_lines") or []:
-            print(f"[smolvla2] {line}", flush=True)
-
-    def _on_shutdown(self) -> None:
-        # Drain any queued action chunks safely.
-        queue = self.state.get("action_queue")
-        if isinstance(queue, deque):
-            queue.clear()
-        print("[smolvla2] runtime stopped", flush=True)

src/lerobot/policies/smolvla2/inference/runtime_state.py

@@ -1,91 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Runtime state passed between inference steps each tick.
-
-The runtime threads a single dict through the pipeline; this module
-documents the shape and provides factories. We use a plain ``dict``
-rather than a frozen dataclass because steps freely add and remove
-keys (``events_this_tick``, ``messages_pending``, ``tool_calls_pending``,
-…) and dataclass field churn would just get in the way.
-
-Stable keys (read by multiple steps):
-
-  task          str             the current top-level task
-  current_plan  str | None      latest plan emitted by the planner
-  current_subtask str | None    latest subtask the policy is executing
-  current_memory str | None     latest compressed memory
-  recent_interjection str | None  most recent user interjection text (consumed)
-
-  action_queue  collections.deque[Tensor]  pending action chunks
-  tool_calls_pending list[dict]  parsed but not-yet-dispatched tool calls
-
-  events_this_tick list[str]    triggers consumed this tick
-  _tick         Tick            current tick (set by the loop)
-
-  log_lines     list[str]       human-readable status lines printed each tick
-"""
-
-from __future__ import annotations
-
-from collections import deque
-from typing import Any
-
-
-def initial_runtime_state(task: str | None = None) -> dict[str, Any]:
-    """Build a fresh runtime state dict with sensible defaults."""
-    return {
-        "task": task,
-        "current_plan": None,
-        "current_subtask": None,
-        "current_memory": None,
-        "recent_interjection": None,
-        "action_queue": deque(),
-        "tool_calls_pending": [],
-        "events_this_tick": [],
-        "log_lines": [],
-        "stop": False,
-    }
-
-
-def take_event(state: dict[str, Any], event_name: str) -> bool:
-    """Pop ``event_name`` from ``events_this_tick`` if present.
-
-    Steps that consume an event call this so the same event doesn't
-    re-fire on a sibling step within the same tick.
-    """
-    events: list[str] = state.get("events_this_tick") or []
-    if event_name in events:
-        events.remove(event_name)
-        return True
-    return False
-
-
-def push_log(state: dict[str, Any], line: str) -> None:
-    """Append ``line`` to the per-tick log buffer; the runtime prints
-    it at the end of the tick."""
-    state.setdefault("log_lines", []).append(line)
-
-
-def set_if_changed(state: dict[str, Any], key: str, value: Any, label: str | None = None) -> bool:
-    """Update ``state[key]`` and log a diff line if the value changed.
-
-    Returns ``True`` if the value actually changed.
-    """
-    prev = state.get(key)
-    if prev == value:
-        return False
-    state[key] = value
-    if label is not None:
-        push_log(state, f"  {label}: {value}")
-    return True

src/lerobot/policies/smolvla2/inference/steps.py

@@ -1,382 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Inference steps for the SmolVLA2 multi-rate runtime.
-
-Each step is a tiny class with a ``trigger`` and an ``__call__(state)``;
-the runtime applies them in order each tick. When a step's trigger
-doesn't fire, the step is a no-op and the runtime moves on.
-
-Stream-to-step mapping mirrors the ``smolvla2_hirobot.yaml`` recipe:
-
-* ``LowLevelForward``        — calls ``policy.select_action`` for the
-                                action chunk; trained by
-                                ``low_level_execution``
-* ``EnqueueChunk``           — pushes the chunk to ``action_queue``
-* ``DispatchAction``         — pops one action per control tick and
-                                forwards to the robot
-* ``HighLevelSubtaskFwd``    — calls ``policy.select_message`` for the
-                                next subtask; trained by
-                                ``high_level_subtask``
-* ``MemoryUpdateFwd``        — fires on subtask boundary; trained by
-                                ``memory_update``
-* ``UserInterjectionFwd``    — fires on stdin interjection; trained by
-                                ``user_interjection_response``
-* ``AskVQAFwd``              — fires on stdin question; trained by
-                                ``ask_vqa_*``
-* ``DispatchToolCalls``      — pops ``tool_calls_pending`` and calls
-                                the matching ``Tool`` instance
-"""
-
-from __future__ import annotations
-
-import logging
-import re
-from dataclasses import dataclass, field
-from typing import Any
-
-from .runtime_state import push_log, set_if_changed, take_event
-from .triggers import EventTrigger, HzTrigger, Trigger
-
-logger = logging.getLogger(__name__)
-
-
-# ---------------------------------------------------------------------------
-# Step base + runner
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class InferenceStep:
-    """A trigger-gated callable. Subclasses override :meth:`run`."""
-
-    trigger: Trigger
-
-    def __call__(self, state: dict[str, Any]) -> dict[str, Any]:
-        if not self.trigger.should_fire(state["_tick"], state):
-            return state
-        return self.run(state) or state
-
-    def run(self, state: dict[str, Any]) -> dict[str, Any] | None:  # pragma: no cover
-        raise NotImplementedError
-
-
-# ---------------------------------------------------------------------------
-# Low-level (action) path
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class LowLevelForward(InferenceStep):
-    """Run the policy's action head and produce one action chunk."""
-
-    policy: Any = None
-    observation_provider: Any = None
-    """Callable ``() -> dict``: returns the current observation batch
-    (already preprocessed). Typically wraps the robot's camera /
-    proprio reads. ``None`` in dry-run mode → step skips."""
-
-    trigger: Trigger = field(default_factory=lambda: HzTrigger(hz=4.0))
-
-    def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
-        if self.policy is None or self.observation_provider is None:
-            return None
-        observation = self.observation_provider()
-        if observation is None:
-            return None
-        action = self.policy.select_action(observation)
-        # SmolVLA returns a single action; if the underlying policy
-        # streams chunks, split per-step here. For v1 we just enqueue
-        # the result.
-        state.setdefault("action_queue", []).append(action)
-        return None
-
-
-@dataclass
-class DispatchAction(InferenceStep):
-    """Pop one action per tick and hand it to the robot.
-
-    In dry-run mode (``robot_executor=None``) the step still pops the
-    queue so it doesn't grow unbounded — the popped tensor is logged
-    instead of executed.
-    """
-
-    robot_executor: Any = None
-    trigger: Trigger = field(default_factory=lambda: HzTrigger(hz=50.0))
-
-    def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
-        queue = state.get("action_queue")
-        if not queue:
-            return None
-        action = queue.popleft() if hasattr(queue, "popleft") else queue.pop(0)
-        if self.robot_executor is not None:
-            self.robot_executor(action)
-        return None
-
-
-# ---------------------------------------------------------------------------
-# High-level (text) paths — all use policy.select_message
-# ---------------------------------------------------------------------------
-
-
-def _build_text_batch(policy: Any, prompt_messages: list[dict[str, Any]]) -> dict[str, Any]:
-    """Tokenize a list of chat messages into the batch shape
-    ``select_message`` expects.
-
-    Lazy fallback: re-uses the policy's preprocessor by piggy-backing
-    on the chat tokenizer step. Production use should construct the
-    batch from a real observation; here we focus on the *language*
-    path which is independent of camera observations.
-    """
-    from transformers import AutoTokenizer  # noqa: PLC0415
-
-    tokenizer = AutoTokenizer.from_pretrained(policy.config.vlm_model_name)
-    if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    text_messages = [_strip_recipe_keys(m) for m in prompt_messages]
-    ids = tokenizer.apply_chat_template(
-        text_messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_tensors="pt",
-    )
-    if isinstance(ids, list):
-        ids = ids[0] if ids else []
-    if hasattr(ids, "ndim") and ids.ndim == 1:
-        ids = ids.unsqueeze(0)
-    attn = (ids != tokenizer.pad_token_id) if tokenizer.pad_token_id is not None else None
-    return {"lang_tokens": ids, "lang_masks": attn, "tokenizer": tokenizer}
-
-
-def _strip_recipe_keys(m: dict[str, Any]) -> dict[str, Any]:
-    new = dict(m)
-    new.pop("stream", None)
-    new.pop("target", None)
-    return new
-
-
-@dataclass
-class HighLevelSubtaskFwd(InferenceStep):
-    """At ~1 Hz, ask the policy for the next subtask."""
-
-    policy: Any = None
-    trigger: Trigger = field(default_factory=lambda: HzTrigger(hz=1.0))
-
-    def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
-        if self.policy is None or not state.get("task"):
-            return None
-        ctx = _control_context_messages(state)
-        msg = _generate_with_policy(self.policy, ctx)
-        if msg:
-            changed = set_if_changed(state, "current_subtask", msg, label="subtask")
-            if changed:
-                # Subtask change is a downstream trigger.
-                state.setdefault("events_this_tick", []).append("subtask_change")
-        return None
-
-
-@dataclass
-class MemoryUpdateFwd(InferenceStep):
-    """On subtask boundary, refresh the compressed memory."""
-
-    policy: Any = None
-    trigger: Trigger = field(default_factory=lambda: EventTrigger("subtask_change"))
-
-    def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
-        # Don't consume the event — multiple steps may want to react.
-        if self.policy is None:
-            return None
-        ctx = _control_context_messages(state, include_completed=True)
-        new_memory = _generate_with_policy(self.policy, ctx)
-        if new_memory:
-            set_if_changed(state, "current_memory", new_memory, label="memory")
-        return None
-
-
-@dataclass
-class UserInterjectionFwd(InferenceStep):
-    """On stdin interjection, refresh the plan + emit a paired ``say``."""
-
-    policy: Any = None
-    trigger: Trigger = field(default_factory=lambda: EventTrigger("user_interjection"))
-
-    def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
-        if self.policy is None or not take_event(state, "user_interjection"):
-            return None
-        ctx = _control_context_messages(
-            state,
-            extra_user=state.get("recent_interjection"),
-        )
-        out = _generate_with_policy(self.policy, ctx)
-        if not out:
-            return None
-        # Heuristic split: model is trained to emit one assistant turn
-        # carrying both plan text AND a `say` tool call. Look for a
-        # "<say>...</say>" or "say(...)" marker; fall back to whole
-        # text → plan, no speech.
-        plan_text, speech_text = _split_plan_and_say(out)
-        if plan_text:
-            set_if_changed(state, "current_plan", plan_text, label="plan")
-        if speech_text:
-            push_log(state, f"  speech: {speech_text}")
-            state.setdefault("tool_calls_pending", []).append(
-                {
-                    "type": "function",
-                    "function": {"name": "say", "arguments": {"text": speech_text}},
-                }
-            )
-            state.setdefault("events_this_tick", []).append("tool_call_pending")
-        # Mark interjection consumed.
-        state["recent_interjection"] = None
-        return None
-
-
-@dataclass
-class AskVQAFwd(InferenceStep):
-    """On stdin question, answer a frame-grounded VQA."""
-
-    policy: Any = None
-    trigger: Trigger = field(default_factory=lambda: EventTrigger("user_vqa_query"))
-
-    def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
-        if self.policy is None or not take_event(state, "user_vqa_query"):
-            return None
-        question = state.get("recent_vqa_query")
-        if not question:
-            return None
-        ctx = _control_context_messages(state, extra_user=question)
-        answer = _generate_with_policy(self.policy, ctx)
-        if answer:
-            push_log(state, f"  vqa: {answer}")
-        state["recent_vqa_query"] = None
-        return None
-
-
-# ---------------------------------------------------------------------------
-# Tool dispatch
-# ---------------------------------------------------------------------------
-
-
-@dataclass
-class DispatchToolCalls(InferenceStep):
-    """Pop ``tool_calls_pending`` and execute them via :data:`TOOL_REGISTRY`."""
-
-    tools: dict[str, Any] = field(default_factory=dict)
-    trigger: Trigger = field(default_factory=lambda: EventTrigger("tool_call_pending"))
-
-    def run(self, state: dict[str, Any]) -> dict[str, Any] | None:
-        take_event(state, "tool_call_pending")
-        pending = state.get("tool_calls_pending") or []
-        for call in pending:
-            try:
-                fn = (call or {}).get("function") or {}
-                name = fn.get("name")
-                args = fn.get("arguments") or {}
-                tool = self.tools.get(name)
-                if tool is None:
-                    push_log(state, f"  [warn] tool {name!r} not registered — skipping call")
-                    continue
-                tool.call(args)
-            except Exception as exc:  # noqa: BLE001
-                push_log(state, f"  [error] tool dispatch failed: {exc}")
-        state["tool_calls_pending"] = []
-        return None
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
-def _control_context_messages(
-    state: dict[str, Any],
-    *,
-    include_completed: bool = False,
-    extra_user: str | None = None,
-) -> list[dict[str, Any]]:
-    """Build a chat-template-ready prompt from current runtime state.
-
-    Mirrors what ``smolvla2_hirobot.yaml`` renders into ``${task}\nPlan:
-    ${plan}\nMemory: ${memory}`` for the high-level branches.
-    """
-    parts: list[str] = []
-    task = state.get("task") or ""
-    parts.append(task)
-    if state.get("current_plan"):
-        parts.append(f"Plan: {state['current_plan']}")
-    if state.get("current_memory"):
-        parts.append(f"Memory: {state['current_memory']}")
-    if include_completed and state.get("current_subtask"):
-        parts.append(f"Completed subtask: {state['current_subtask']}")
-    head = "\n".join(parts)
-    msgs: list[dict[str, Any]] = [{"role": "user", "content": head}]
-    if extra_user:
-        msgs.append({"role": "user", "content": extra_user})
-    return msgs
-
-
-def _generate_with_policy(policy: Any, messages: list[dict[str, Any]]) -> str:
-    """Drive ``policy.select_message`` with a minimal text-only batch.
-
-    Best-effort: the runtime today doesn't construct a full
-    observation batch with images / state for text generation; the
-    text-head was trained over images + lang + state, so generations
-    here may differ in distribution from training. This is acceptable
-    for a v1 REPL; a follow-up will plug in the real observation.
-    """
-    if not hasattr(policy, "select_message"):
-        return ""
-    text_batch = _build_text_batch(policy, messages)
-    # ``select_message`` expects a real batch with OBS_LANGUAGE_TOKENS.
-    # The minimal text-only batch we build doesn't have images / state,
-    # so we either run a text-only forward (handled by SmolVLA2 when
-    # supported) or skip and return empty. v1 returns empty when the
-    # policy can't handle it; the runtime logs and continues.
-    try:
-        # Convert to the OBS_LANGUAGE_TOKENS / OBS_LANGUAGE_ATTENTION_MASK
-        # keys ``select_message`` uses internally.
-        from lerobot.utils.constants import (  # noqa: PLC0415
-            OBS_LANGUAGE_ATTENTION_MASK,
-            OBS_LANGUAGE_TOKENS,
-        )
-
-        batch = {
-            OBS_LANGUAGE_TOKENS: text_batch["lang_tokens"],
-            OBS_LANGUAGE_ATTENTION_MASK: text_batch["lang_masks"],
-        }
-        return policy.select_message(batch, tokenizer=text_batch["tokenizer"])
-    except Exception as exc:  # noqa: BLE001
-        logger.debug("select_message fell back: %s", exc)
-        return ""
-
-
-_SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTALL)
-
-
-def _split_plan_and_say(text: str) -> tuple[str, str]:
-    """Pull a ``<say>...</say>`` snippet out of ``text``; remainder is plan.
-
-    The training-time tool-call serializer wraps ``say(text="…")`` in a
-    deterministic textual marker so prefix-LM-style training learns to
-    emit it. The runtime parses it back here. If no marker is present,
-    the entire text is treated as plan with no speech.
-    """
-    if not text:
-        return "", ""
-    match = _SAY_RE.search(text)
-    if not match:
-        return text.strip(), ""
-    speech = match.group(1).strip().strip('"').strip("'")
-    plan = (text[: match.start()] + text[match.end() :]).strip()
-    return plan, speech

src/lerobot/policies/smolvla2/inference/triggers.py

@@ -1,117 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Trigger primitives for SmolVLA2's multi-rate inference runtime.
-
-Mirrors the plan's Section "Runtime orchestration": each
-``InferenceStep`` is gated by a :class:`Trigger` that decides per tick
-whether the step fires. Two trigger flavours cover all the cadences
-the canonical recipe needs:
-
-* :class:`HzTrigger` for periodic beats (action chunks at ~3-5 Hz,
-  high-level subtask generation at ~1 Hz, action dispatch at ~50 Hz)
-* :class:`EventTrigger` for one-shot reactions (subtask boundary →
-  memory update; user interjection → plan refresh; user VQA query →
-  vqa answer; pending tool call → dispatcher)
-
-Triggers are stateless except for ``HzTrigger``'s last-fire timestamp.
-The runtime stores the :class:`Tick` clock as ``state["_tick"]`` so
-every step shares a single time source.
-"""
-
-from __future__ import annotations
-
-import time
-from dataclasses import dataclass, field
-from typing import Any, Protocol
-
-
-@dataclass
-class Tick:
-    """Single tick from :class:`TickClock`. Carries time references the
-    runtime steps consume to gate themselves."""
-
-    index: int
-    """Monotonic counter — increments by one per tick."""
-
-    monotonic_seconds: float
-    """``time.monotonic()`` at the start of this tick."""
-
-
-@dataclass
-class TickClock:
-    """Drives the runtime loop at up to ``max_rate_hz``.
-
-    Sleeps just enough between :meth:`advance` calls to enforce the
-    rate. With ``max_rate_hz=50`` the loop wakes ~every 20ms; the
-    higher-level ``HzTrigger`` slices that timeline into sub-cadences.
-    """
-
-    max_rate_hz: float = 50.0
-    _index: int = field(default=0, init=False)
-    _last_seconds: float | None = field(default=None, init=False)
-
-    def advance(self) -> Tick:
-        period = 1.0 / max(self.max_rate_hz, 0.1)
-        now = time.monotonic()
-        if self._last_seconds is not None:
-            sleep_for = (self._last_seconds + period) - now
-            if sleep_for > 0:
-                time.sleep(sleep_for)
-                now = time.monotonic()
-        self._last_seconds = now
-        self._index += 1
-        return Tick(index=self._index, monotonic_seconds=now)
-
-
-class Trigger(Protocol):
-    """Decide whether the next ``InferenceStep`` should fire."""
-
-    def should_fire(self, tick: Tick, state: dict[str, Any]) -> bool: ...
-
-
-@dataclass
-class HzTrigger:
-    """Fire at most ``hz`` times per second."""
-
-    hz: float
-    _last_seconds: float | None = field(default=None, init=False)
-
-    def should_fire(self, tick: Tick, state: dict[str, Any]) -> bool:
-        period = 1.0 / max(self.hz, 1e-6)
-        if self._last_seconds is None or (tick.monotonic_seconds - self._last_seconds) >= period:
-            self._last_seconds = tick.monotonic_seconds
-            return True
-        return False
-
-
-@dataclass
-class EventTrigger:
-    """Fire when ``event_name`` is in ``state["events_this_tick"]``.
-
-    The runtime fills ``events_this_tick`` once per tick from:
-
-    * stdin / network input (``user_interjection``, ``user_vqa_query``,
-      ``stop``)
-    * internal state transitions (``subtask_change``,
-      ``tool_call_pending``)
-
-    The list is consumed (cleared at the end of the tick) so events
-    fire at most once.
-    """
-
-    event_name: str
-
-    def should_fire(self, tick: Tick, state: dict[str, Any]) -> bool:
-        events: list[str] = state.get("events_this_tick") or []
-        return self.event_name in events

src/lerobot/policies/smolvla2/modeling_smolvla2.py

@@ -1,387 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""SmolVLA2 modeling — dual-head subclass of SmolVLAPolicy.
-
-Adds:
-
-* an unfrozen SmolVLM ``lm_head`` so language tokens can be supervised,
-* a forward path that runs the flow head, the text head, or both,
-  driven by ``batch["predict_actions"]`` and ``batch["text_labels"]``
-  produced by :class:`SmolVLA2ChatTokenizerStep` (the previous commit on
-  this branch).
-
-Per-sample routing — within one batch:
-
-* ``predict_actions[i] = True`` ⇒ sample ``i`` contributes to the flow
-  loss (action chunk supervision).
-* ``predict_actions[i] = False`` ⇒ sample ``i`` is masked out of the
-  flow loss; only its text tokens (where ``text_labels[i, t] != -100``)
-  contribute to the LM-head cross-entropy.
-
-Falls back to ``SmolVLAPolicy.forward`` cleanly when neither
-``text_labels`` nor ``predict_actions`` is in the batch — unannotated
-datasets keep working unchanged.
-"""
-
-from __future__ import annotations
-
-import math
-from typing import Any
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor
-
-from lerobot.utils.constants import (
-    ACTION,
-    OBS_LANGUAGE_ATTENTION_MASK,
-    OBS_LANGUAGE_TOKENS,
-    OBS_STATE,
-)
-
-from ..smolvla.modeling_smolvla import SmolVLAPolicy, make_att_2d_masks
-from .configuration_smolvla2 import SmolVLA2Config
-
-
-class SmolVLA2Policy(SmolVLAPolicy):
-    """SmolVLA + re-enabled SmolVLM language head."""
-
-    config_class = SmolVLA2Config
-    name = "smolvla2"
-
-    def __init__(self, config: SmolVLA2Config, dataset_stats: dict[str, dict[str, Tensor]] | None = None):
-        if not isinstance(config, SmolVLA2Config):
-            config = SmolVLA2Config(
-                **{
-                    f.name: getattr(config, f.name)
-                    for f in config.__dataclass_fields__.values()
-                    if hasattr(config, f.name)
-                }
-            )
-        super().__init__(config, dataset_stats=dataset_stats)
-        if config.unfreeze_lm_head and config.text_loss_weight > 0:
-            self._unfreeze_lm_head()
-
-    # ------------------------------------------------------------------
-    # Backbone surgery
-    # ------------------------------------------------------------------
-
-    def _unfreeze_lm_head(self) -> None:
-        """Re-enable gradients on the SmolVLM ``lm_head`` (and the bits
-        of the text path SmolVLA freezes) so the text-loss can flow back.
-        """
-        vlm_with_expert = getattr(self.model, "vlm_with_expert", None)
-        if vlm_with_expert is None:
-            return
-        vlm = getattr(vlm_with_expert, "vlm", None)
-        if vlm is None:
-            return
-        for name, param in vlm.named_parameters():
-            if "lm_head" in name or "text_model.model.norm.weight" in name:
-                param.requires_grad = True
-
-    # ------------------------------------------------------------------
-    # Forward
-    # ------------------------------------------------------------------
-
-    def forward(
-        self,
-        batch: dict[str, Tensor],
-        noise: Tensor | None = None,
-        time: Tensor | None = None,
-        reduction: str = "mean",
-    ) -> tuple[Tensor, dict[str, Any]]:
-        """Forward pass with optional dual-head loss.
-
-        Two routing knobs from the batch (produced by
-        :class:`SmolVLA2ChatTokenizerStep`):
-
-        * ``text_labels`` — per-token labels with ``-100`` for non-target
-          positions. Triggers the text-loss path through ``lm_head``.
-        * ``predict_actions`` — per-sample bool tensor. ``True`` ⇒
-          include this sample's action chunk in the flow loss.
-
-        When neither is present, delegate to ``SmolVLAPolicy.forward``.
-        """
-        text_labels = batch.get("text_labels")
-        predict_actions_t = batch.get("predict_actions")
-
-        has_text_data = (
-            text_labels is not None
-            and isinstance(text_labels, Tensor)
-            and self.config.text_loss_weight > 0
-        )
-        has_per_sample_routing = (
-            predict_actions_t is not None and isinstance(predict_actions_t, Tensor)
-        )
-
-        if not has_text_data and not has_per_sample_routing:
-            return super().forward(batch, noise=noise, time=time, reduction=reduction)
-
-        loss_dict: dict[str, Any] = {}
-        device = batch[OBS_STATE].device
-        total = torch.zeros((), device=device, dtype=torch.float32)
-
-        # ------------------------------------------------------------
-        # Flow loss path — only when at least one sample wants actions.
-        # ------------------------------------------------------------
-        run_flow = self.config.flow_loss_weight > 0 and (
-            not has_per_sample_routing or bool(predict_actions_t.any().item())
-        )
-        if run_flow and ACTION in batch:
-            per_sample_flow, flow_diag = super().forward(
-                batch, noise=noise, time=time, reduction="none"
-            )
-            # ``per_sample_flow`` has shape (B,) from the SmolVLA
-            # reduction="none" branch.
-            if has_per_sample_routing:
-                mask = predict_actions_t.to(per_sample_flow.dtype)
-                masked = per_sample_flow * mask
-                denom = mask.sum().clamp(min=1.0)
-                flow_loss = masked.sum() / denom
-            else:
-                flow_loss = per_sample_flow.mean()
-            total = total + self.config.flow_loss_weight * flow_loss
-            loss_dict["flow_loss"] = float(flow_loss.detach().item())
-            for k, v in flow_diag.items():
-                loss_dict[f"flow_{k}"] = v
-
-        # ------------------------------------------------------------
-        # Text loss path — prefix-only forward → lm_head → CE.
-        # ------------------------------------------------------------
-        if has_text_data:
-            text_loss = self._compute_text_loss(batch, text_labels)
-            total = total + self.config.text_loss_weight * text_loss
-            loss_dict["text_loss"] = float(text_loss.detach().item())
-
-        loss_dict["loss"] = float(total.detach().item())
-
-        if reduction == "none":
-            # Per-sample loss isn't meaningfully defined for the dual
-            # path; broadcast the scalar to (B,) for caller compat.
-            return total.expand(batch[OBS_STATE].shape[0]), loss_dict
-        return total, loss_dict
-
-    # ------------------------------------------------------------------
-    # Text-loss internals
-    # ------------------------------------------------------------------
-
-    def _compute_text_loss(self, batch: dict[str, Tensor], text_labels: Tensor) -> Tensor:
-        """Cross-entropy on the SmolVLM ``lm_head`` over target tokens."""
-        if self.config.adapt_to_pi_aloha:
-            batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
-
-        images, img_masks = self.prepare_images(batch)
-        state = self.prepare_state(batch)
-        lang_tokens = batch[OBS_LANGUAGE_TOKENS]
-        lang_masks = batch[OBS_LANGUAGE_ATTENTION_MASK]
-
-        prefix_embs, prefix_pad_masks, prefix_att_masks = self.model.embed_prefix(
-            images, img_masks, lang_tokens, lang_masks, state=state
-        )
-        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
-        prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
-
-        # Prefix-only forward.
-        out_pair, _ = self.model.vlm_with_expert.forward(
-            attention_mask=prefix_att_2d_masks,
-            position_ids=prefix_position_ids,
-            past_key_values=None,
-            inputs_embeds=[prefix_embs, None],
-            use_cache=False,
-            fill_kv_cache=False,
-        )
-        prefix_out = out_pair[0] if isinstance(out_pair, (tuple, list)) else out_pair
-        if prefix_out is None:
-            raise RuntimeError(
-                "SmolVLA2: vlm_with_expert.forward returned no prefix hidden "
-                "states — text-loss path needs them."
-            )
-
-        # Lang token positions inside the prefix. ``embed_prefix`` lays
-        # out the prefix as ``[image_blocks..., lang, state]`` so the
-        # lang range is identifiable from the trailing state size and
-        # the known lang length.
-        num_lang = lang_tokens.shape[1]
-        state_for_dim = state if state.ndim >= 2 else state[:, None]
-        num_state = state_for_dim.shape[1] if state_for_dim.ndim >= 2 else 1
-        if num_state < 1:
-            num_state = 1
-        prefix_len = prefix_out.shape[1]
-        lang_end = prefix_len - num_state
-        lang_start = lang_end - num_lang
-        if lang_start < 0 or lang_end > prefix_len:
-            raise RuntimeError(
-                f"SmolVLA2: could not locate lang token range in prefix "
-                f"(prefix_len={prefix_len}, num_lang={num_lang}, "
-                f"num_state={num_state})."
-            )
-
-        lang_hidden = prefix_out[:, lang_start:lang_end]
-        vlm = self.model.vlm_with_expert.vlm
-        logits = vlm.lm_head(lang_hidden)  # (B, num_lang, vocab)
-
-        if text_labels.shape[1] != num_lang:
-            common = min(text_labels.shape[1], num_lang)
-            logits = logits[:, :common]
-            text_labels = text_labels[:, :common]
-
-        # Standard next-token CE: hidden state at position t predicts
-        # token at position t+1. Shift logits left, labels right by 1.
-        # Without this, the loss is identity-mapped and the LM head
-        # learns nothing useful — see HuggingFace ``LlamaForCausalLM``
-        # for the same convention.
-        shift_logits = logits[:, :-1, :].contiguous()
-        shift_labels = text_labels[:, 1:].contiguous().long()
-        loss = F.cross_entropy(
-            shift_logits.reshape(-1, shift_logits.shape[-1]),
-            shift_labels.reshape(-1),
-            ignore_index=-100,
-        )
-        return loss
-
-    # ------------------------------------------------------------------
-    # Inference: text generation
-    # ------------------------------------------------------------------
-
-    @torch.no_grad()
-    def select_message(
-        self,
-        batch: dict[str, Tensor],
-        *,
-        max_new_tokens: int = 256,
-        eos_token_id: int | None = None,
-        temperature: float = 0.0,
-        top_p: float = 1.0,
-        tokenizer: Any = None,
-    ) -> str:
-        """Generate text continuation from the chat-templated prompt.
-
-        AR decoding with KV caching reused from SmolVLA's inference
-        path. Batch size is assumed to be 1 (the runtime calls this
-        per-event). Returns the decoded string of new tokens (the
-        prompt itself is not included).
-
-        Parameters
-        ----------
-        batch:
-            Already through the SmolVLA2 preprocessor — expects
-            ``OBS_IMAGES_*``, ``OBS_STATE``, ``OBS_LANGUAGE_TOKENS``,
-            ``OBS_LANGUAGE_ATTENTION_MASK``.
-        max_new_tokens:
-            Hard cap on generated tokens; stops earlier on EOS.
-        eos_token_id:
-            Override the tokenizer's EOS. ``None`` ⇒ use the
-            tokenizer's default.
-        temperature, top_p:
-            ``temperature=0`` does greedy argmax (default — matches
-            training distribution most closely). Set ``temperature>0``
-            with optional ``top_p<1`` for nucleus sampling.
-        tokenizer:
-            Optional pre-loaded tokenizer to avoid the cold-start
-            ``AutoTokenizer.from_pretrained`` round-trip on every call.
-        """
-        self.eval()
-
-        if tokenizer is None:
-            from transformers import AutoTokenizer  # noqa: PLC0415
-
-            tokenizer = AutoTokenizer.from_pretrained(self.config.vlm_model_name)
-        if eos_token_id is None:
-            eos_token_id = tokenizer.eos_token_id
-
-        images, img_masks = self.prepare_images(batch)
-        state = self.prepare_state(batch)
-        lang_tokens = batch[OBS_LANGUAGE_TOKENS]
-        lang_masks = batch[OBS_LANGUAGE_ATTENTION_MASK]
-
-        # 1) Embed prefix (images + lang + state) and run with KV cache.
-        prefix_embs, prefix_pad_masks, prefix_att_masks = self.model.embed_prefix(
-            images, img_masks, lang_tokens, lang_masks, state=state
-        )
-        prefix_2d = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
-        prefix_pos = torch.cumsum(prefix_pad_masks, dim=1) - 1
-        out_pair, past_kv = self.model.vlm_with_expert.forward(
-            attention_mask=prefix_2d,
-            position_ids=prefix_pos,
-            past_key_values=None,
-            inputs_embeds=[prefix_embs, None],
-            use_cache=True,
-            fill_kv_cache=True,
-        )
-        prefix_out = out_pair[0] if isinstance(out_pair, (tuple, list)) else out_pair
-        if prefix_out is None:
-            raise RuntimeError("select_message: prefix forward returned no hidden states.")
-
-        vlm = self.model.vlm_with_expert.vlm
-
-        # 2) Initial logits — sample first new token from the last
-        #    prefix position.
-        last_hidden = prefix_out[:, -1:]
-        device = last_hidden.device
-        bsize = prefix_embs.shape[0]
-        cur_pos = int(prefix_embs.shape[1])
-
-        generated: list[int] = []
-        for _ in range(max_new_tokens):
-            logits_step = vlm.lm_head(last_hidden)[:, -1]  # (B, V)
-            next_ids = self._sample_next_token(logits_step, temperature, top_p)
-            tok_id = int(next_ids[0].item())
-            generated.append(tok_id)
-            if eos_token_id is not None and tok_id == eos_token_id:
-                break
-
-            # 3) Embed the new token and forward with KV cache.
-            new_emb = self.model.vlm_with_expert.embed_language_tokens(
-                next_ids.unsqueeze(0)
-            )
-            new_emb = new_emb * math.sqrt(new_emb.shape[-1])
-
-            new_pos = torch.full((bsize, 1), cur_pos, device=device, dtype=torch.long)
-            new_attn = torch.ones((bsize, cur_pos + 1), device=device, dtype=torch.bool)
-
-            out_pair, past_kv = self.model.vlm_with_expert.forward(
-                attention_mask=new_attn,
-                position_ids=new_pos,
-                past_key_values=past_kv,
-                inputs_embeds=[new_emb, None],
-                use_cache=True,
-                fill_kv_cache=True,
-            )
-            new_prefix_out = out_pair[0] if isinstance(out_pair, (tuple, list)) else out_pair
-            last_hidden = new_prefix_out[:, -1:]
-            cur_pos += 1
-
-        return tokenizer.decode(generated, skip_special_tokens=True).strip()
-
-    @staticmethod
-    def _sample_next_token(
-        logits: Tensor, temperature: float, top_p: float
-    ) -> Tensor:
-        """Pick one token id per batch row from ``logits``."""
-        if temperature <= 0.0:
-            return logits.argmax(dim=-1)
-        scaled = logits / max(temperature, 1e-6)
-        probs = F.softmax(scaled, dim=-1)
-        if top_p < 1.0:
-            sorted_probs, sorted_idx = probs.sort(dim=-1, descending=True)
-            cum = sorted_probs.cumsum(dim=-1)
-            mask = cum > top_p
-            # Always keep the most-likely token.
-            mask[..., 0] = False
-            sorted_probs = sorted_probs.masked_fill(mask, 0.0)
-            sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True).clamp(min=1e-9)
-            pick = torch.multinomial(sorted_probs, num_samples=1)
-            return sorted_idx.gather(-1, pick).squeeze(-1)
-        return torch.multinomial(probs, num_samples=1).squeeze(-1)

src/lerobot/policies/smolvla2/processor_smolvla2.py

@@ -1,131 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""SmolVLA2 processor pipelines.
-
-When ``config.recipe_path`` is set, the pre-processor pipeline becomes:
-
-    rename observations
-    add batch dim
-    RenderMessagesStep(recipe)              # PR 1: language_*  → messages
-    SmolVLA2ChatTokenizerStep(...)          # chat template + label mask + predict_actions
-    DeviceProcessorStep
-    NormalizerProcessorStep
-
-When ``config.recipe_path`` is ``None``, we delegate to SmolVLA's
-plain task-string pipeline so unannotated datasets still work.
-
-Post-processor is unchanged from SmolVLA.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Any
-
-import torch
-
-from lerobot.configs.recipe import TrainingRecipe
-from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    RenderMessagesStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
-)
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-
-from ..smolvla.processor_smolvla import make_smolvla_pre_post_processors
-from .chat_processor_smolvla2 import SmolVLA2ChatTokenizerStep
-from .configuration_smolvla2 import SmolVLA2Config
-
-
-def make_smolvla2_pre_post_processors(
-    config: SmolVLA2Config,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """Build SmolVLA2's pre/post-processor pipelines.
-
-    With ``recipe_path`` set, inserts the recipe-rendering step and the
-    chat-template tokenizer that emits ``text_labels`` and
-    ``predict_actions`` for the dual-loss path. Without it, falls back
-    to SmolVLA's plain task-string pipeline so unannotated datasets
-    keep working unchanged.
-    """
-    if not config.recipe_path:
-        return make_smolvla_pre_post_processors(config, dataset_stats=dataset_stats)
-
-    recipe = _load_recipe(config.recipe_path)
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        RenderMessagesStep(recipe=recipe),
-        SmolVLA2ChatTokenizerStep(
-            tokenizer_name=config.vlm_model_name,
-            max_length=config.tokenizer_max_length,
-            padding=config.pad_language_to,
-        ),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features,
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        DeviceProcessorStep(device="cpu"),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )
-
-
-def _load_recipe(path_str: str) -> TrainingRecipe:
-    """Resolve ``path_str`` to a ``TrainingRecipe``.
-
-    Accepts an absolute path or a path relative to
-    ``src/lerobot/configs/`` so recipe authors can write
-    ``--policy.recipe_path=recipes/smolvla2_hirobot.yaml``.
-    """
-    p = Path(path_str)
-    if not p.is_absolute() and not p.exists():
-        from lerobot.configs import recipe as _recipe_module  # noqa: PLC0415
-
-        configs_dir = Path(_recipe_module.__file__).resolve().parent
-        candidate = configs_dir / path_str
-        if candidate.exists():
-            p = candidate
-    return TrainingRecipe.from_yaml(p)

src/lerobot/processor/__init__.py

@@ -93,7 +93,6 @@ from .relative_action_processor import (
     to_relative_actions,
 )
 from .rename_processor import RenameObservationsProcessorStep, rename_stats
-from .render_messages_processor import RenderMessagesStep
 from .tokenizer_processor import ActionTokenizerProcessorStep, TokenizerProcessorStep
 
 __all__ = [
@@ -129,7 +128,6 @@ __all__ = [
     "make_default_robot_observation_processor",
     "AbsoluteActionsProcessorStep",
     "RelativeActionsProcessorStep",
-    "RenderMessagesStep",
     "MapDeltaActionToRobotActionStep",
     "MapTensorToDeltaActionDictStep",
     "NewLineTaskProcessorStep",