refactor(rewards): migrate SARM from policies/sarm/ to rewards/sarm/

src/lerobot/rewards/sarm/compute_rabc_weights.py

@@ -0,0 +1,870 @@
+#!/usr/bin/env python
+
+# 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.
+
+"""
+Compute SARM progress values for RA-BC (Reward-Aware Behavior Cloning) weighting.
+
+This script processes all frames in a dataset with SARM to compute progress values [0, 1].
+The results are saved as a parquet file that can be loaded during training for RA-BC weighting.
+
+Uses multi-output extraction: each SARM query returns progress for 9 frames, so we only
+need ~num_frames/30 queries instead of one per frame (~30x speedup).
+
+Usage:
+    # Full RA-BC computation with visualizations
+    python src/lerobot/rewards/sarm/compute_rabc_weights.py \\
+        --dataset-repo-id lerobot/aloha_sim_insertion_human \\
+        --reward-model-path <USER>/sarm_single_uni4
+
+    # Faster computation with stride (compute every 5 frames, interpolate the rest)
+    python src/lerobot/rewards/sarm/compute_rabc_weights.py \\
+        --dataset-repo-id lerobot/aloha_sim_insertion_human \\
+        --reward-model-path <USER>/sarm_single_uni4 \\
+        --stride 5
+
+    # Visualize predictions only (no RA-BC computation)
+    python src/lerobot/rewards/sarm/compute_rabc_weights.py \\
+        --dataset-repo-id lerobot/aloha_sim_insertion_human \\
+        --reward-model-path <USER>/sarm_single_uni4 \\
+        --visualize-only \\
+        --num-visualizations 5
+
+The output is saved to the dataset's local cache directory as 'sarm_progress.parquet'.
+"""
+
+import argparse
+import logging
+from pathlib import Path
+
+import matplotlib.gridspec as gridspec
+import matplotlib.pyplot as plt
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import torch
+from tqdm import tqdm
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.rewards.sarm.modeling_sarm import SARMRewardModel
+from lerobot.rewards.sarm.processor_sarm import make_sarm_pre_post_processors
+from lerobot.rewards.sarm.sarm_utils import normalize_stage_tau
+
+
+def get_reward_model_path_from_parquet(parquet_path: Path) -> str | None:
+    """Read reward_model_path from parquet metadata if available."""
+    if not parquet_path.exists():
+        return None
+    try:
+        metadata = pq.read_metadata(parquet_path).schema.to_arrow_schema().metadata
+        if metadata and b"reward_model_path" in metadata:
+            return metadata[b"reward_model_path"].decode()
+    except Exception:  # nosec B110
+        return None
+    return None
+
+
+def load_sarm_resources(
+    dataset_repo_id: str,
+    reward_model_path: str,
+    device: str = "cuda",
+) -> tuple[LeRobotDataset, SARMRewardModel, any]:
+    """
+    Load SARM model, dataset, and preprocessor.
+
+    Returns:
+        Tuple of (dataset, reward_model, preprocessor)
+    """
+    logging.info(f"Loading model: {reward_model_path}")
+    reward_model = SARMRewardModel.from_pretrained(reward_model_path)
+    reward_model.config.device = device
+    reward_model.to(device).eval()
+
+    image_key = reward_model.config.image_key
+    state_key = reward_model.config.state_key
+    delta_indices = reward_model.config.observation_delta_indices
+
+    logging.info(f"Loading dataset: {dataset_repo_id}")
+    temp_dataset = LeRobotDataset(dataset_repo_id, download_videos=True)
+    fps = temp_dataset.fps
+
+    delta_timestamps = {
+        image_key: [idx / fps for idx in delta_indices],
+        state_key: [idx / fps for idx in delta_indices],
+    }
+    dataset = LeRobotDataset(dataset_repo_id, delta_timestamps=delta_timestamps)
+    logging.info(f"Dataset: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
+
+    preprocess, _ = make_sarm_pre_post_processors(
+        config=reward_model.config,
+        dataset_stats=dataset.meta.stats,
+        dataset_meta=dataset.meta,
+    )
+
+    return dataset, reward_model, preprocess
+
+
+def to_numpy_image(img) -> np.ndarray:
+    """Convert image tensor to numpy uint8 (H, W, C)."""
+    if isinstance(img, torch.Tensor):
+        img = img.cpu().numpy()
+    if img.ndim == 4:
+        # Take center frame for bidirectional sampling
+        img = img[img.shape[0] // 2]
+    if img.shape[0] in [1, 3]:
+        img = np.transpose(img, (1, 2, 0))
+    if img.dtype != np.uint8:
+        # Handle normalized images (may have negative values or values > 1)
+        img = img.astype(np.float32)
+        img = (img - img.min()) / (img.max() - img.min() + 1e-8)  # Normalize to [0, 1]
+        img = (img * 255).astype(np.uint8)
+    return img
+
+
+def visualize_episode(
+    frames, progress_preds, stage_preds, title, output_path, stage_labels, gt_progress=None, gt_stages=None
+):
+    """Create visualization with progress plot, stage probabilities, and sample frames.
+
+    Same as sarm_inference_visualization.py
+    """
+    num_stages = stage_preds.shape[1]
+    colors = plt.cm.tab10(np.linspace(0, 1, num_stages))
+    frame_indices = np.arange(len(progress_preds))
+
+    fig = plt.figure(figsize=(14, 12))
+    gs = gridspec.GridSpec(3, 1, height_ratios=[2, 1, 1], hspace=0.3)
+    ax_progress, ax_stages, ax_frames = fig.add_subplot(gs[0]), fig.add_subplot(gs[1]), fig.add_subplot(gs[2])
+
+    # Progress plot
+    ax_progress.plot(frame_indices, progress_preds, linewidth=2, color="#2E86AB", label="Predicted")
+    ax_progress.fill_between(frame_indices, 0, progress_preds, alpha=0.3, color="#2E86AB")
+    if gt_progress is not None:
+        ax_progress.plot(
+            frame_indices, gt_progress, linewidth=2, color="#28A745", linestyle="--", label="Ground Truth"
+        )
+    ax_progress.axhline(y=1.0, color="gray", linestyle="--", alpha=0.5)
+    ax_progress.set_ylabel("Progress")
+    ax_progress.set_title(f'Task: "{title}"', fontweight="bold")
+    ax_progress.set_ylim(-0.05, 1.1)
+    ax_progress.legend(loc="upper left")
+    ax_progress.grid(True, alpha=0.3)
+
+    # Stage predictions
+    ax_stages.stackplot(
+        frame_indices,
+        *[stage_preds[:, i] for i in range(num_stages)],
+        colors=colors,
+        alpha=0.8,
+        labels=stage_labels,
+    )
+    if gt_stages is not None:
+        for change_idx in np.where(np.diff(gt_stages) != 0)[0] + 1:
+            ax_stages.axvline(x=change_idx, color="black", linestyle="-", alpha=0.7, linewidth=1.5)
+    ax_stages.set_xlabel("Frame")
+    ax_stages.set_ylabel("Stage Probability")
+    ax_stages.set_ylim(0, 1)
+    ax_stages.legend(loc="upper left", ncol=min(num_stages, 5), fontsize=8)
+    ax_stages.grid(True, alpha=0.3)
+
+    # Sample frames
+    ax_frames.axis("off")
+    num_sample = 8
+    sample_indices = np.linspace(0, len(frames) - 1, num_sample, dtype=int)
+    h, w = frames[0].shape[:2]
+    combined = np.zeros((h, w * num_sample, 3), dtype=np.uint8)
+    for i, idx in enumerate(sample_indices):
+        frame = frames[idx]
+        if frame.shape[-1] == 1:
+            frame = np.repeat(frame, 3, axis=-1)
+        combined[:, i * w : (i + 1) * w] = frame
+        stage_name = stage_labels[np.argmax(stage_preds[idx])][:12]
+        ax_frames.text(
+            i * w + w / 2,
+            -10,
+            f"Frame {idx}\n{progress_preds[idx]:.2f}\n{stage_name}",
+            ha="center",
+            va="top",
+            fontsize=7,
+        )
+    ax_frames.imshow(combined)
+    ax_frames.set_title("Sample Frames", pad=20)
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"Saved: {output_path}")
+
+
+def visualize_sarm_predictions(
+    dataset: LeRobotDataset,
+    reward_model: SARMRewardModel,
+    preprocess,
+    episode_indices: list[int],
+    head_mode: str,
+    output_dir: Path,
+    num_display_frames: int = 5,
+    stride: int = 1,
+):
+    """
+    Visualize SARM predictions for multiple episodes.
+
+    Computes predictions for every frame by default. With stride > 1, computes predictions
+    every N frames and interpolates (progress + stage probabilities) for visualization.
+
+    Args:
+        dataset: LeRobotDataset with delta_timestamps configured
+        reward_model: Loaded SARM model
+        preprocess: Preprocessor from make_sarm_pre_post_processors
+        episode_indices: List of episode indices to visualize
+        head_mode: "sparse", "dense", or "both"
+        output_dir: Directory to save visualizations
+        num_display_frames: Number of frames to display in thumbnail strip (default: 5)
+        stride: Compute predictions every N frames, interpolate the rest (default: 1)
+    """
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    image_key = reward_model.config.image_key
+    state_key = reward_model.config.state_key
+    dual_mode = reward_model.config.uses_dual_heads
+    device = reward_model.device
+
+    # Center frame index for bidirectional sampling
+    target_idx = reward_model.config.n_obs_steps // 2
+
+    # Determine which heads to visualize
+    schemes_to_viz = []
+    if head_mode in ("sparse", "both") or not dual_mode:
+        schemes_to_viz.append("sparse")
+    if head_mode in ("dense", "both") and dual_mode:
+        schemes_to_viz.append("dense")
+
+    # Set preprocessor to eval mode to disable augmentations
+    if hasattr(preprocess, "eval"):
+        preprocess.eval()
+    for step in preprocess.steps:
+        if hasattr(step, "eval"):
+            step.eval()
+
+    for episode_idx in episode_indices:
+        ep = dataset.meta.episodes[episode_idx]
+        ep_start = ep["dataset_from_index"]
+        ep_end = ep["dataset_to_index"]
+        task = dataset[ep_start].get("task", "perform the task")
+        num_frames = ep_end - ep_start
+
+        # Select frames for display thumbnails (evenly sampled from begin to end)
+        display_indices = set(
+            [
+                ep_start + int(i * (num_frames - 1) / (num_display_frames - 1))
+                for i in range(num_display_frames)
+            ]
+            if num_frames >= num_display_frames
+            else list(range(ep_start, ep_end))
+        )
+        viz_frames = {}
+
+        # Load display frames up-front (stride mode might skip them otherwise).
+        for frame_idx in display_indices:
+            sample = dataset[frame_idx]
+            viz_frames[frame_idx] = to_numpy_image(sample[image_key])
+
+        # Initialize storage for each scheme
+        scheme_data = {}
+        for scheme in schemes_to_viz:
+            num_stages = getattr(reward_model.config, f"num_{scheme}_stages")
+            scheme_data[scheme] = {
+                "viz_progress": np.full(num_frames, np.nan),
+                "viz_stages": np.full((num_frames, num_stages), np.nan),
+                "viz_gt_progress": np.full(num_frames, np.nan),
+                "viz_gt_stages": np.full(num_frames, np.nan),
+                "target_key": f"{scheme}_targets",
+                "num_stages": num_stages,
+                "temporal_props": getattr(reward_model.config, f"{scheme}_temporal_proportions"),
+                "subtask_names": getattr(reward_model.config, f"{scheme}_subtask_names"),
+            }
+
+        if stride > 1:
+            logging.info(f"Visualization stride={stride}: inferring every {stride} frames and interpolating")
+
+        # Process frames one at a time to avoid memory buildup
+        frame_indices = list(range(ep_start, ep_end, stride))
+        if (ep_end - 1) not in frame_indices:
+            frame_indices.append(ep_end - 1)
+        frame_indices = sorted(set(frame_indices))
+
+        for frame_idx in tqdm(frame_indices, desc=f"Episode {episode_idx}", leave=False):
+            local_idx = frame_idx - ep_start
+            sample = dataset[frame_idx]
+
+            batch = {
+                image_key: sample[image_key],
+                "task": task,
+                "index": frame_idx,
+                "episode_index": episode_idx,
+            }
+            if state_key in sample:
+                batch[state_key] = sample[state_key]
+
+            with torch.no_grad():
+                processed = preprocess(batch)
+                video_features = processed["video_features"].to(device)
+                text_features = processed["text_features"].to(device)
+                state_features = processed.get("state_features")
+                if state_features is not None:
+                    state_features = state_features.to(device)
+                lengths = processed.get("lengths")
+
+                for scheme in schemes_to_viz:
+                    sd = scheme_data[scheme]
+
+                    # Ground truth
+                    # In stride visualization mode, ground-truth plots can be misleading
+                    # (only sparse points are available), so we skip GT.
+                    if stride == 1 and sd["target_key"] in processed:
+                        gt_target = processed[sd["target_key"]][0, target_idx].cpu().item()
+                        sd["viz_gt_stages"][local_idx] = int(gt_target)
+                        sd["viz_gt_progress"][local_idx] = normalize_stage_tau(
+                            gt_target,
+                            num_stages=sd["num_stages"],
+                            temporal_proportions=sd["temporal_props"],
+                            subtask_names=sd["subtask_names"],
+                        )
+
+                    # Predictions
+                    reward, stage_probs = reward_model.calculate_rewards(
+                        text_embeddings=text_features,
+                        video_embeddings=video_features,
+                        state_features=state_features,
+                        lengths=lengths,
+                        return_all_frames=True,
+                        return_stages=True,
+                        head_mode=scheme,
+                    )
+
+                    # Handle both tensor and numpy outputs
+                    if isinstance(reward, torch.Tensor):
+                        reward = reward.cpu().numpy()
+                        stage_probs = stage_probs.cpu().numpy()
+
+                    if reward.ndim == 2:
+                        sd["viz_progress"][local_idx] = reward[0, target_idx]
+                        sd["viz_stages"][local_idx] = stage_probs[0, target_idx, :]
+                    else:
+                        sd["viz_progress"][local_idx] = reward[target_idx]
+                        sd["viz_stages"][local_idx] = stage_probs[target_idx, :]
+
+                # Clear GPU memory after each frame
+                del processed, video_features, text_features
+                if state_features is not None:
+                    del state_features
+
+            torch.cuda.empty_cache()
+
+        # Interpolate predictions back to per-frame arrays for smooth visualization.
+        if stride > 1:
+            all_local = np.arange(num_frames)
+            for scheme in schemes_to_viz:
+                sd = scheme_data[scheme]
+
+                valid = np.isfinite(sd["viz_progress"])
+                valid_idx = np.where(valid)[0]
+                if valid_idx.size >= 1:
+                    sd["viz_progress"] = interpolate_progress(
+                        valid_idx, sd["viz_progress"][valid_idx], all_local
+                    )
+
+                    stage_interp = np.zeros_like(sd["viz_stages"], dtype=np.float32)
+                    for s in range(sd["num_stages"]):
+                        stage_interp[:, s] = interpolate_progress(
+                            valid_idx, sd["viz_stages"][valid_idx, s], all_local
+                        )
+
+                    stage_interp = np.clip(stage_interp, 0.0, 1.0)
+                    row_sums = stage_interp.sum(axis=1, keepdims=True)
+                    nz = row_sums.squeeze(-1) > 0
+                    stage_interp[nz] = stage_interp[nz] / row_sums[nz]
+                    sd["viz_stages"] = stage_interp
+                else:
+                    # No valid points: keep NaNs/zeros; visualization will be empty.
+                    sd["viz_stages"] = np.nan_to_num(sd["viz_stages"], nan=0.0)
+
+        # Generate visualization for each head
+        ordered_viz_frames = [viz_frames[idx] for idx in sorted(display_indices)]
+        for scheme in schemes_to_viz:
+            sd = scheme_data[scheme]
+            stage_labels = sd["subtask_names"] or [f"Stage {i + 1}" for i in range(sd["num_stages"])]
+            viz_path = output_dir / f"sarm_prediction_ep{episode_idx}_{scheme}.png"
+
+            visualize_episode(
+                frames=np.array(ordered_viz_frames),
+                progress_preds=sd["viz_progress"],
+                stage_preds=sd["viz_stages"],
+                title=f"{task} (Episode {episode_idx})",
+                output_path=viz_path,
+                stage_labels=stage_labels,
+                gt_progress=sd["viz_gt_progress"] if not np.all(np.isnan(sd["viz_gt_progress"])) else None,
+                gt_stages=sd["viz_gt_stages"] if not np.all(np.isnan(sd["viz_gt_stages"])) else None,
+            )
+
+        # Clear memory between episodes
+        torch.cuda.empty_cache()
+
+    logging.info(f"Visualizations saved to: {output_dir.absolute()}")
+
+
+def generate_all_frame_indices(ep_start: int, ep_end: int, frame_gap: int = 30) -> list[int]:
+    """Generate all frame indices, ordered by offset for cache-friendly access.
+
+    Orders frames as: [0, 30, 60...], [1, 31, 61...], ..., [29, 59, 89...]
+    This groups frames that share similar temporal windows together.
+    """
+    num_frames = ep_end - ep_start
+    indices = []
+    for offset in range(frame_gap):
+        for frame_rel in range(offset, num_frames, frame_gap):
+            indices.append(ep_start + frame_rel)
+    return indices
+
+
+def interpolate_progress(
+    computed_indices: np.ndarray,
+    computed_values: np.ndarray,
+    all_indices: np.ndarray,
+) -> np.ndarray:
+    """Linearly interpolate values to fill in gaps (robust to NaNs / edge cases)."""
+    computed_indices = np.asarray(computed_indices)
+    computed_values = np.asarray(computed_values)
+    all_indices = np.asarray(all_indices)
+
+    mask = np.isfinite(computed_values)
+    if mask.sum() == 0:
+        return np.full(all_indices.shape, np.nan, dtype=np.float32)
+    if mask.sum() == 1:
+        return np.full(all_indices.shape, float(computed_values[mask][0]), dtype=np.float32)
+
+    out = np.interp(all_indices, computed_indices[mask], computed_values[mask])
+    return out.astype(np.float32)
+
+
+def compute_sarm_progress(
+    dataset_repo_id: str,
+    reward_model_path: str,
+    output_path: str | None = None,
+    head_mode: str = "sparse",
+    device: str = "cuda",
+    num_visualizations: int = 5,
+    output_dir: str = "./sarm_viz",
+    stride: int = 1,
+):
+    """
+    Compute SARM progress predictions for all frames in a dataset.
+
+    Args:
+        dataset_repo_id: HuggingFace dataset repo ID or local path
+        reward_model_path: Path to pretrained SARM model
+        output_path: Path to save results. If None, saves to dataset's cache directory
+        head_mode: SARM head to use ("sparse", "dense", or "both")
+        device: Device to use for inference
+        num_visualizations: Number of episodes to visualize (0 to skip)
+        output_dir: Directory to save visualizations
+        stride: Compute progress every N frames, interpolate the rest (default: 1 = every frame)
+    """
+    dataset, reward_model, preprocess = load_sarm_resources(dataset_repo_id, reward_model_path, device)
+
+    # Set preprocessor to eval mode to disable augmentations
+    if hasattr(preprocess, "eval"):
+        preprocess.eval()
+    for step in preprocess.steps:
+        if hasattr(step, "eval"):
+            step.eval()
+
+    image_key = reward_model.config.image_key
+    state_key = reward_model.config.state_key
+    frame_gap = reward_model.config.frame_gap
+    num_episodes = dataset.num_episodes
+    total_frames = dataset.num_frames
+    logging.info(f"Processing {total_frames} frames across {num_episodes} episodes")
+
+    # Determine which heads to compute
+    dual_mode = reward_model.config.uses_dual_heads
+    compute_sparse = head_mode in ("sparse", "both") or not dual_mode
+    compute_dense = head_mode in ("dense", "both") and dual_mode
+
+    # Storage arrays
+    all_indices = []
+    all_episode_indices = []
+    all_frame_indices = []
+    all_progress_sparse = [] if compute_sparse else None
+    all_progress_dense = [] if compute_dense else None
+
+    if stride > 1:
+        logging.info(f"Using stride={stride}: computing every {stride} frames, interpolating the rest")
+
+    # Process all episodes
+    for episode_idx in tqdm(range(num_episodes), desc="Episodes"):
+        ep = dataset.meta.episodes[episode_idx]
+        ep_start = ep["dataset_from_index"]
+        ep_end = ep["dataset_to_index"]
+
+        # Get task description
+        task = dataset[ep_start].get("task", "perform the task")
+
+        # Generate frames to compute (with stride applied)
+        all_ep_indices = generate_all_frame_indices(ep_start, ep_end, frame_gap)
+        if stride > 1:
+            # Only compute every stride-th frame (relative to episode start)
+            compute_indices = [idx for idx in all_ep_indices if (idx - ep_start) % stride == 0]
+            # Always include last frame for better interpolation at episode end
+            last_frame = ep_end - 1
+            if last_frame not in compute_indices:
+                compute_indices.append(last_frame)
+            compute_indices = sorted(set(compute_indices))
+        else:
+            compute_indices = all_ep_indices
+
+        center_idx = reward_model.config.n_obs_steps // 2  # Center of bidirectional window
+
+        # Dictionary to collect results
+        frame_results = {}
+
+        for query_idx in tqdm(compute_indices, desc=f"  Ep {episode_idx}", leave=False):
+            try:
+                sample = dataset[query_idx]
+
+                batch = {
+                    image_key: sample[image_key],
+                    "task": task,
+                    "index": query_idx,
+                    "episode_index": episode_idx,
+                }
+                if state_key in sample:
+                    batch[state_key] = sample[state_key]
+
+                with torch.no_grad():
+                    processed = preprocess(batch)
+                    video_features = processed["video_features"].to(device)
+                    text_features = processed["text_features"].to(device)
+                    state_features = processed.get("state_features")
+                    if state_features is not None:
+                        state_features = state_features.to(device)
+                    lengths = processed.get("lengths")
+
+                    sparse_val = np.nan
+                    dense_val = np.nan
+
+                    # Compute sparse prediction for center frame
+                    if compute_sparse:
+                        sparse_progress = reward_model.calculate_rewards(
+                            text_embeddings=text_features,
+                            video_embeddings=video_features,
+                            state_features=state_features,
+                            lengths=lengths,
+                            return_all_frames=True,
+                            head_mode="sparse",
+                        )
+                        sparse_val = float(
+                            sparse_progress[0, center_idx]
+                            if sparse_progress.ndim == 2
+                            else sparse_progress[center_idx]
+                        )
+
+                    # Compute dense prediction for center frame
+                    if compute_dense:
+                        dense_progress = reward_model.calculate_rewards(
+                            text_embeddings=text_features,
+                            video_embeddings=video_features,
+                            state_features=state_features,
+                            lengths=lengths,
+                            return_all_frames=True,
+                            head_mode="dense",
+                        )
+                        dense_val = float(
+                            dense_progress[0, center_idx]
+                            if dense_progress.ndim == 2
+                            else dense_progress[center_idx]
+                        )
+
+                    frame_results[query_idx] = (sparse_val, dense_val)
+
+            except Exception as e:
+                logging.warning(f"Failed to process frame {query_idx}: {e}")
+
+        # Interpolate to get values for all frames
+        computed_indices = np.array(sorted(frame_results.keys()))
+        computed_sparse = (
+            np.array([frame_results[i][0] for i in computed_indices]) if compute_sparse else None
+        )
+        computed_dense = np.array([frame_results[i][1] for i in computed_indices]) if compute_dense else None
+
+        # All frame indices for this episode
+        all_frame_idx_array = np.arange(ep_start, ep_end)
+
+        if stride > 1 and len(computed_indices) > 1:
+            # Interpolate progress values
+            if compute_sparse:
+                interp_sparse = interpolate_progress(computed_indices, computed_sparse, all_frame_idx_array)
+            if compute_dense:
+                interp_dense = interpolate_progress(computed_indices, computed_dense, all_frame_idx_array)
+        else:
+            # No interpolation needed
+            interp_sparse = computed_sparse if compute_sparse else None
+            interp_dense = computed_dense if compute_dense else None
+
+        # Store results for all frames
+        for i, frame_idx in enumerate(all_frame_idx_array):
+            local_idx = frame_idx - ep_start
+            all_indices.append(frame_idx)
+            all_episode_indices.append(episode_idx)
+            all_frame_indices.append(local_idx)
+            if compute_sparse:
+                if stride > 1 and len(computed_indices) > 1:
+                    all_progress_sparse.append(float(interp_sparse[i]))
+                elif frame_idx in frame_results:
+                    all_progress_sparse.append(frame_results[frame_idx][0])
+                else:
+                    all_progress_sparse.append(np.nan)
+            if compute_dense:
+                if stride > 1 and len(computed_indices) > 1:
+                    all_progress_dense.append(float(interp_dense[i]))
+                elif frame_idx in frame_results:
+                    all_progress_dense.append(frame_results[frame_idx][1])
+                else:
+                    all_progress_dense.append(np.nan)
+
+    # Create output table
+    table_data = {
+        "index": np.array(all_indices, dtype=np.int64),
+        "episode_index": np.array(all_episode_indices, dtype=np.int64),
+        "frame_index": np.array(all_frame_indices, dtype=np.int64),
+    }
+    if compute_sparse:
+        table_data["progress_sparse"] = np.array(all_progress_sparse, dtype=np.float32)
+    if compute_dense:
+        table_data["progress_dense"] = np.array(all_progress_dense, dtype=np.float32)
+
+    # Sort by index
+    df = pa.table(table_data).to_pandas()
+    df = df.sort_values("index").reset_index(drop=True)
+    final_table = pa.Table.from_pandas(df, preserve_index=False)
+
+    # Add metadata with reward model path
+    metadata = {b"reward_model_path": reward_model_path.encode()}
+    final_table = final_table.replace_schema_metadata(metadata)
+
+    # Determine output path
+    output_path = Path(dataset.root) / "sarm_progress.parquet" if output_path is None else Path(output_path)
+
+    # Save
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    pq.write_table(final_table, output_path)
+    logging.info(f"Saved {len(final_table)} frame progress values to {output_path}")
+
+    # Print statistics
+    if "progress_sparse" in df.columns:
+        valid = df["progress_sparse"].dropna()
+        logging.info(
+            f"Sparse progress: mean={valid.mean():.4f}, std={valid.std():.4f}, "
+            f"min={valid.min():.4f}, max={valid.max():.4f}"
+        )
+
+    if "progress_dense" in df.columns:
+        valid = df["progress_dense"].dropna()
+        logging.info(
+            f"Dense progress: mean={valid.mean():.4f}, std={valid.std():.4f}, "
+            f"min={valid.min():.4f}, max={valid.max():.4f}"
+        )
+
+    # Visualize episodes after processing
+    if num_visualizations > 0:
+        viz_episodes = list(range(min(num_visualizations, num_episodes)))
+        logging.info(f"Generating {len(viz_episodes)} visualizations...")
+        visualize_sarm_predictions(
+            dataset=dataset,
+            reward_model=reward_model,
+            preprocess=preprocess,
+            episode_indices=viz_episodes,
+            head_mode=head_mode,
+            output_dir=Path(output_dir),
+            stride=stride,
+        )
+
+    return output_path
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Compute SARM progress values for RA-BC weighting or visualize SARM predictions",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # Full RA-BC computation with visualizations
+    python src/lerobot/rewards/sarm/compute_rabc_weights.py \\
+        --dataset-repo-id lerobot/aloha_sim_insertion_human \\
+        --reward-model-path <USER>/sarm_single_uni4
+
+    # Visualize predictions only (no RA-BC computation)
+    python src/lerobot/rewards/sarm/compute_rabc_weights.py \\
+        --dataset-repo-id lerobot/aloha_sim_insertion_human \\
+        --reward-model-path <USER>/sarm_single_uni4 \\
+        --visualize-only \\
+        --num-visualizations 10
+        """,
+    )
+    parser.add_argument(
+        "--dataset-repo-id",
+        type=str,
+        required=True,
+        help="HuggingFace dataset repo ID or local path",
+    )
+    parser.add_argument(
+        "--reward-model-path",
+        type=str,
+        default=None,
+        help="Path to pretrained SARM model (reads from existing parquet metadata if not provided)",
+    )
+    parser.add_argument(
+        "--output-path",
+        type=str,
+        default=None,
+        help="Output path for parquet. If not set, saves to dataset's cache directory",
+    )
+    parser.add_argument(
+        "--head-mode",
+        type=str,
+        default="sparse",
+        choices=["sparse", "dense", "both"],
+        help="SARM head to use (default: sparse)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        help="Device to use (default: cuda)",
+    )
+    # Visualization options
+    parser.add_argument(
+        "--visualize-only",
+        action="store_true",
+        help="Only visualize SARM predictions (no RA-BC computation)",
+    )
+    parser.add_argument(
+        "--num-visualizations",
+        type=int,
+        default=5,
+        help="Number of episodes to visualize (default: 5, set to 0 to skip)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="./sarm_viz",
+        help="Output directory for visualizations (default: ./sarm_viz)",
+    )
+    parser.add_argument(
+        "--push-to-hub",
+        action="store_true",
+        help="Upload progress file to the dataset repo on HuggingFace Hub",
+        default=True,
+    )
+    parser.add_argument(
+        "--stride",
+        type=int,
+        default=1,
+        help="Compute progress every N frames, interpolate the rest (default: 1 = every frame)",
+    )
+
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+
+    # Try to get reward_model_path from parquet metadata if not provided
+    reward_model_path = args.reward_model_path
+    if reward_model_path is None:
+        # Load dataset to find parquet path
+        temp_dataset = LeRobotDataset(args.dataset_repo_id, download_videos=False)
+        parquet_path = Path(temp_dataset.root) / "sarm_progress.parquet"
+        reward_model_path = get_reward_model_path_from_parquet(parquet_path)
+        if reward_model_path:
+            logging.info(f"Using reward model from parquet metadata: {reward_model_path}")
+        else:
+            raise ValueError(
+                "--reward-model-path is required (no existing parquet with model metadata found)"
+            )
+
+    # Handle visualize-only mode
+    if args.visualize_only:
+        dataset, reward_model, preprocess = load_sarm_resources(
+            args.dataset_repo_id, reward_model_path, args.device
+        )
+        logging.info(f"Visualization-only mode: visualizing {args.num_visualizations} episodes")
+        viz_episodes = list(range(min(args.num_visualizations, dataset.num_episodes)))
+        visualize_sarm_predictions(
+            dataset=dataset,
+            reward_model=reward_model,
+            preprocess=preprocess,
+            episode_indices=viz_episodes,
+            head_mode=args.head_mode,
+            output_dir=Path(args.output_dir),
+            stride=args.stride,
+        )
+        print(f"\nVisualizations saved to: {Path(args.output_dir).absolute()}")
+        return
+
+    # Full RABC computation (compute_sarm_progress loads model/dataset itself)
+    output_path = compute_sarm_progress(
+        dataset_repo_id=args.dataset_repo_id,
+        reward_model_path=reward_model_path,
+        output_path=args.output_path,
+        head_mode=args.head_mode,
+        device=args.device,
+        num_visualizations=args.num_visualizations,
+        output_dir=args.output_dir,
+        stride=args.stride,
+    )
+
+    print(f"\nSARM progress values saved to: {output_path}")
+
+    # Upload to Hub if requested
+    if args.push_to_hub:
+        from huggingface_hub import HfApi
+
+        api = HfApi()
+        hub_path = "sarm_progress.parquet"
+
+        print(f"\nUploading to Hub: {args.dataset_repo_id}/{hub_path}")
+        api.upload_file(
+            path_or_fileobj=str(output_path),
+            path_in_repo=hub_path,
+            repo_id=args.dataset_repo_id,
+            repo_type="dataset",
+        )
+        print(
+            f"Successfully uploaded to: https://huggingface.co/datasets/{args.dataset_repo_id}/blob/main/{hub_path}"
+        )
+
+        print("\nTo use in training, add to your config:")
+        print("  use_rabc: true")
+        print(f"  rabc_progress_path: hf://datasets/{args.dataset_repo_id}/{hub_path}")
+        print("  rabc_head_mode: sparse  # or dense")
+    else:
+        print("\nTo use in training, add to your config:")
+        print("  use_rabc: true")
+        print(f"  rabc_progress_path: {output_path}")
+        print("  rabc_head_mode: sparse  # or dense")
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/rewards/sarm/configuration_sarm.py

@@ -0,0 +1,232 @@
+#!/usr/bin/env python
+
+# Copyright 2025 Qianzhong Chen, Justin Yu, Mac Schwager, Pieter Abbeel, Yide Shentu, Philipp Wu
+# and 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.
+
+"""
+SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation.
+Paper: https://arxiv.org/abs/2509.25358
+"""
+
+from dataclasses import dataclass, field
+
+from lerobot.configs.rewards import RewardModelConfig
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
+from lerobot.utils.constants import OBS_IMAGES, OBS_STATE
+
+
+@RewardModelConfig.register_subclass("sarm")
+@dataclass
+class SARMConfig(RewardModelConfig):
+    """Configuration class for SARM (Stage-Aware Reward Modeling).
+
+    Supports three annotation modes:
+
+    1. single_stage (default): No annotations needed. Uses the episode's task description
+       as a single stage covering the entire episode.
+
+    2. dense_only: Uses dense (fine-grained) annotations from VLM, with an auto-generated
+       single sparse "task" stage covering the full episode. The dense head learns detailed
+       subtask progression while sparse provides overall task completion.
+
+    3. dual: Full dual-head mode with both sparse (high-level) and dense (fine-grained)
+       annotations from VLM. Both heads are trained on their respective annotations.
+
+    The annotation_mode determines how sparse_temporal_proportions and dense_temporal_proportions
+    are loaded/generated during model initialization.
+    """
+
+    annotation_mode: str = "single_stage"  # "single_stage", "dense_only", or "dual"
+    n_obs_steps: int = 8  # Number of observation history steps
+    frame_gap: int = 30  # Frame gap between frames (at 30 fps = 1 second)
+    max_rewind_steps: int = 4  # Maximum rewind steps for temporal augmentation
+
+    # Architecture params
+    image_dim: int = 512
+    text_dim: int = 512
+    hidden_dim: int = 768
+    num_heads: int = 12
+    num_layers: int = 8
+    max_state_dim: int = 32
+    drop_n_last_frames: int = 1
+    batch_size: int = 64
+    clip_batch_size: int = 64
+    dropout: float = 0.1
+    stage_loss_weight: float = 1.0
+
+    rewind_probability: float = 0.8
+    language_perturbation_probability: float = 0.2
+
+    # Sparse annotations (high-level stages)
+    num_sparse_stages: int = 1
+    sparse_subtask_names: list | None = None
+    sparse_temporal_proportions: list | None = None
+
+    # Dense annotations (fine-grained stages)
+    num_dense_stages: int | None = None
+    dense_subtask_names: list | None = None
+    dense_temporal_proportions: list | None = None
+
+    pretrained_model_path: str | None = None
+    image_key: str = OBS_IMAGES + ".top"
+    state_key: str = OBS_STATE
+
+    # Populated by the processor (video_features, state_features, text_features)
+    input_features: dict = field(default_factory=lambda: {})
+
+    # Output features (updated in __post_init__)
+    output_features: dict = field(
+        default_factory=lambda: {
+            "stage": PolicyFeature(shape=(9, 5), type=FeatureType.REWARD),
+            "progress": PolicyFeature(shape=(9, 1), type=FeatureType.REWARD),
+        }
+    )
+
+    normalization_mapping: dict[str, NormalizationMode] = field(
+        default_factory=lambda: {
+            "VISUAL": NormalizationMode.IDENTITY,
+            "STATE": NormalizationMode.MEAN_STD,
+            "LANGUAGE": NormalizationMode.IDENTITY,
+            "REWARD": NormalizationMode.IDENTITY,
+        }
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        if self.annotation_mode not in ["single_stage", "dense_only", "dual"]:
+            raise ValueError(
+                f"annotation_mode must be 'single_stage', 'dense_only', or 'dual', got {self.annotation_mode}"
+            )
+
+        if self.annotation_mode == "single_stage":
+            self.num_sparse_stages = 1
+            self.sparse_subtask_names = ["task"]
+            self.sparse_temporal_proportions = [1.0]
+            self.num_dense_stages = None
+            self.dense_subtask_names = None
+            self.dense_temporal_proportions = None
+
+        elif self.annotation_mode == "dense_only":
+            self.num_sparse_stages = 1
+            self.sparse_subtask_names = ["task"]
+            self.sparse_temporal_proportions = [1.0]
+
+        self.input_features = {}
+        self.output_features = {}
+
+        if self.image_key:
+            self.input_features[self.image_key] = PolicyFeature(shape=(480, 640, 3), type=FeatureType.VISUAL)
+
+        self.input_features[self.state_key] = PolicyFeature(
+            shape=(self.max_state_dim,),
+            type=FeatureType.STATE,
+        )
+
+        # Update output features based on annotation_mode
+        if self.annotation_mode in ["dense_only", "dual"]:
+            self.output_features["sparse_stage"] = PolicyFeature(
+                shape=(self.num_frames, self.num_sparse_stages), type=FeatureType.REWARD
+            )
+            self.output_features["sparse_progress"] = PolicyFeature(
+                shape=(self.num_frames, 1), type=FeatureType.REWARD
+            )
+            dense_stages = self.num_dense_stages or self.num_sparse_stages
+            self.output_features["dense_stage"] = PolicyFeature(
+                shape=(self.num_frames, dense_stages), type=FeatureType.REWARD
+            )
+            self.output_features["dense_progress"] = PolicyFeature(
+                shape=(self.num_frames, 1), type=FeatureType.REWARD
+            )
+        else:
+            self.output_features["sparse_stage"] = PolicyFeature(
+                shape=(self.num_frames, self.num_sparse_stages), type=FeatureType.REWARD
+            )
+            self.output_features["sparse_progress"] = PolicyFeature(
+                shape=(self.num_frames, 1), type=FeatureType.REWARD
+            )
+
+        if self.max_rewind_steps >= self.n_obs_steps:
+            raise ValueError(
+                f"max_rewind_steps ({self.max_rewind_steps}) must be less than n_obs_steps ({self.n_obs_steps})"
+            )
+        if self.num_sparse_stages < 1:
+            raise ValueError(f"num_sparse_stages must be at least 1, got {self.num_sparse_stages}")
+        if (
+            self.annotation_mode in ["dense_only", "dual"]
+            and self.num_dense_stages is not None
+            and self.num_dense_stages < 2
+        ):
+            raise ValueError(f"num_dense_stages must be at least 2, got {self.num_dense_stages}")
+
+    def get_optimizer_preset(self) -> AdamWConfig:
+        """Get default optimizer configuration for SARM training."""
+        return AdamWConfig(
+            lr=5e-5,
+            weight_decay=1e-3,
+            betas=(0.9, 0.999),
+            eps=1e-8,
+        )
+
+    def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig:
+        """Get default learning rate scheduler configuration."""
+        return CosineDecayWithWarmupSchedulerConfig(
+            peak_lr=5e-5,
+            decay_lr=5e-6,
+            num_warmup_steps=500,
+            num_decay_steps=50000,
+        )
+
+    def validate_features(self) -> None:
+        pass
+
+    @property
+    def uses_dual_heads(self) -> bool:
+        """Whether the model uses dual heads (dense_only or dual annotation modes)."""
+        return self.annotation_mode in ["dense_only", "dual"]
+
+    @property
+    def num_frames(self) -> int:
+        """Total number of frames in sequence."""
+        return 1 + self.n_obs_steps + self.max_rewind_steps
+
+    @property
+    def max_length(self) -> int:
+        return self.num_frames
+
+    @property
+    def observation_delta_indices(self) -> list[int]:
+        """Bidirectional frame sampling centered on target frame."""
+        half_steps = self.n_obs_steps // 2
+
+        past_deltas = [-self.frame_gap * i for i in range(half_steps, 0, -1)]
+        future_deltas = [self.frame_gap * i for i in range(1, half_steps + 1)]
+        obs_deltas = past_deltas + [0] + future_deltas
+
+        # Rewind placeholders
+        rewind_deltas = [-self.frame_gap * (i + 1) for i in range(self.max_rewind_steps)]
+
+        return obs_deltas + rewind_deltas
+
+    @property
+    def action_delta_indices(self) -> None:
+        """SARM is a reward model, not an action policy."""
+        return None
+
+    @property
+    def reward_delta_indices(self) -> None:
+        return None

src/lerobot/rewards/sarm/modeling_sarm.py

@@ -0,0 +1,627 @@
+#!/usr/bin/env python
+
+# Copyright 2025 Qianzhong Chen, Justin Yu, Mac Schwager, Pieter Abbeel, Yide Shentu, Philipp Wu
+# and 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.
+
+"""
+SARM: Stage-Aware Reward Modeling for Long Horizon Robot Manipulation.
+
+Paper: https://arxiv.org/abs/2509.25358
+
+- StageTransformer: Predicts stage classification (sparse/dense)
+- SubtaskTransformer: Predicts within-stage progress (tau) conditioned on stage
+"""
+
+import json
+import logging
+import random
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F  # noqa: N812
+from torch import Tensor
+
+from lerobot.rewards.pretrained import PreTrainedRewardModel
+from lerobot.rewards.sarm.configuration_sarm import SARMConfig
+from lerobot.rewards.sarm.sarm_utils import (
+    normalize_stage_tau,
+    pad_state_to_max_dim,
+)
+from lerobot.utils.constants import OBS_STR
+
+
+class StageTransformer(nn.Module):
+    """
+    Stage classification transformer for SARM.
+
+    Predicts which stage/subtask the current frame belongs to.
+    Supports both sparse (high-level) and dense (fine-grained) annotation schemes.
+
+    Input streams: [vis_proj, lang_proj, state_proj] concatenated -> (B, N+2, T, D)
+    Output: stage logits (B, T, num_classes)
+    """
+
+    def __init__(
+        self,
+        d_model: int = 512,
+        vis_emb_dim: int = 512,
+        text_emb_dim: int = 512,
+        state_dim: int = 32,
+        n_layers: int = 6,
+        n_heads: int = 8,
+        dropout: float = 0.1,
+        num_cameras: int = 1,
+        num_classes_sparse: int = 4,
+        num_classes_dense: int = 8,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.num_cameras = num_cameras
+
+        # Projections
+        self.lang_proj = nn.Linear(text_emb_dim, d_model)
+        self.visual_proj = nn.Linear(vis_emb_dim, d_model)
+        self.state_proj = nn.Linear(state_dim, d_model)
+
+        # Encoder
+        enc_layer = nn.TransformerEncoderLayer(d_model, n_heads, 4 * d_model, dropout, batch_first=True)
+        self.transformer = nn.TransformerEncoder(enc_layer, n_layers)
+
+        # Positional bias on first visual frame
+        self.first_pos = nn.Parameter(torch.zeros(1, d_model))
+
+        # Shared fusion MLP
+        fused_in = d_model * (num_cameras + 2)
+        self.fusion_backbone = nn.Sequential(
+            nn.LayerNorm(fused_in),
+            nn.Linear(fused_in, d_model),
+            nn.ReLU(),
+        )
+
+        # Scheme-specific heads
+        self.heads = nn.ModuleDict(
+            {
+                "sparse": nn.Linear(d_model, num_classes_sparse),
+                "dense": nn.Linear(d_model, num_classes_dense),
+            }
+        )
+
+    def _prep_lang(self, lang_emb: torch.Tensor, B: int, T: int, D: int) -> torch.Tensor:  # noqa: N803
+        """Prepare language embeddings for fusion."""
+        if lang_emb.dim() == 3:
+            lang_proj = self.lang_proj(lang_emb).unsqueeze(1)
+        else:
+            lang_proj = self.lang_proj(lang_emb).unsqueeze(1).unsqueeze(2).expand(B, 1, T, D)
+        return lang_proj
+
+    def forward(
+        self,
+        img_seq: torch.Tensor,
+        lang_emb: torch.Tensor,
+        state: torch.Tensor,
+        lengths: torch.Tensor,
+        scheme: str = "sparse",
+    ) -> torch.Tensor:
+        assert scheme in self.heads, f"Unknown scheme '{scheme}'. Use one of {list(self.heads.keys())}."
+
+        B, N, T, _ = img_seq.shape  # noqa: N806
+        D = self.d_model  # noqa: N806
+        device = img_seq.device
+
+        vis_proj = self.visual_proj(img_seq)
+        state_proj = self.state_proj(state).unsqueeze(1)
+        lang_proj = self._prep_lang(lang_emb, B, T, D)
+
+        x = torch.cat([vis_proj, lang_proj, state_proj], dim=1)
+        x[:, :N, 0, :] = x[:, :N, 0, :] + self.first_pos
+
+        x_tokens = x.view(B, (N + 2) * T, D)
+        L = x_tokens.size(1)  # noqa: N806
+
+        base_mask = torch.arange(T, device=device).expand(B, T) >= lengths.unsqueeze(1)
+        mask = base_mask.unsqueeze(1).expand(B, N + 2, T).reshape(B, (N + 2) * T)
+
+        causal_mask = torch.triu(torch.ones(L, L, device=device, dtype=torch.bool), diagonal=1)
+
+        h = self.transformer(x_tokens, mask=causal_mask, src_key_padding_mask=mask, is_causal=True)
+
+        h = h.view(B, N + 2, T, D).permute(0, 2, 1, 3).reshape(B, T, (N + 2) * D)
+        fused = self.fusion_backbone(h)
+
+        logits = self.heads[scheme](fused)
+        return logits
+
+
+class SubtaskTransformer(nn.Module):
+    """
+    Subtask progress regression transformer for SARM.
+
+    Predicts within-stage normalized progress (tau) conditioned on stage prior.
+    """
+
+    def __init__(
+        self,
+        d_model: int = 512,
+        vis_emb_dim: int = 512,
+        text_emb_dim: int = 512,
+        state_dim: int = 32,
+        n_layers: int = 6,
+        n_heads: int = 8,
+        dropout: float = 0.1,
+        num_cameras: int = 1,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.num_cameras = num_cameras
+
+        self.lang_proj = nn.Linear(text_emb_dim, d_model)
+        self.visual_proj = nn.Linear(vis_emb_dim, d_model)
+        self.state_proj = nn.Linear(state_dim, d_model)
+
+        enc = nn.TransformerEncoderLayer(d_model, n_heads, 4 * d_model, dropout, batch_first=True)
+        self.transformer = nn.TransformerEncoder(enc, n_layers)
+
+        self.first_pos = nn.Parameter(torch.zeros(1, d_model))
+
+        fused_in = d_model * (num_cameras + 3)
+        self.fusion_backbone = nn.Sequential(
+            nn.LayerNorm(fused_in),
+            nn.Linear(fused_in, d_model),
+            nn.ReLU(),
+        )
+
+        self.heads = nn.ModuleDict(
+            {
+                "sparse": nn.Linear(d_model, 1),
+                "dense": nn.Linear(d_model, 1),
+            }
+        )
+
+    def _prep_lang(self, lang_emb: torch.Tensor, B: int, T: int, D: int) -> torch.Tensor:  # noqa: N803
+        if lang_emb.dim() == 3:
+            return self.lang_proj(lang_emb).unsqueeze(1)
+        else:
+            return self.lang_proj(lang_emb).unsqueeze(1).unsqueeze(2).expand(B, 1, T, D)
+
+    def _stage_to_dmodel(self, stage_prior: torch.Tensor) -> torch.Tensor:
+        B, one, T, C = stage_prior.shape  # noqa: N806
+        D = self.d_model  # noqa: N806
+        if D == C:
+            return stage_prior
+        elif D > C:
+            pad = torch.zeros(B, one, T, D - C, device=stage_prior.device, dtype=stage_prior.dtype)
+            return torch.cat([stage_prior, pad], dim=-1)
+        else:
+            return stage_prior[..., :D]
+
+    def forward(
+        self,
+        img_seq: torch.Tensor,
+        lang_emb: torch.Tensor,
+        state: torch.Tensor,
+        lengths: torch.Tensor,
+        stage_prior: torch.Tensor,
+        scheme: str = "sparse",
+    ) -> torch.Tensor:
+        assert scheme in self.heads, f"Unknown scheme '{scheme}'. Use one of {list(self.heads.keys())}."
+
+        B, N, T, _ = img_seq.shape  # noqa: N806
+        D = self.d_model  # noqa: N806
+        device = img_seq.device
+
+        vis_proj = self.visual_proj(img_seq)
+        state_proj = self.state_proj(state).unsqueeze(1)
+        lang_proj = self._prep_lang(lang_emb, B, T, D)
+        stage_emb = self._stage_to_dmodel(stage_prior)
+
+        x = torch.cat([vis_proj, lang_proj, state_proj, stage_emb], dim=1)
+        x[:, :N, 0, :] = x[:, :N, 0, :] + self.first_pos
+
+        x_tokens = x.view(B, (N + 3) * T, D)
+        L = x_tokens.size(1)  # noqa: N806
+
+        base_mask = torch.arange(T, device=device).expand(B, T) >= lengths.unsqueeze(1)
+        mask = base_mask.unsqueeze(1).expand(B, N + 3, T).reshape(B, (N + 3) * T)
+
+        causal_mask = torch.triu(torch.ones(L, L, device=device, dtype=torch.bool), diagonal=1)
+
+        h = self.transformer(x_tokens, mask=causal_mask, src_key_padding_mask=mask, is_causal=True)
+
+        h = h.view(B, N + 3, T, D)
+        h_flat = h.permute(0, 2, 1, 3).reshape(B, T, (N + 3) * D)
+        fused = self.fusion_backbone(h_flat)
+
+        r = torch.sigmoid(self.heads[scheme](fused)).squeeze(-1)
+        return r
+
+
+def gen_stage_emb(num_classes: int, targets: torch.Tensor) -> torch.Tensor:
+    """Generate one-hot stage embeddings from targets."""
+    idx = targets.long().clamp(min=0, max=num_classes - 1)
+    C = num_classes  # noqa: N806
+    stage_onehot = torch.eye(C, device=targets.device)[idx]
+    stage_onehot = stage_onehot.unsqueeze(1)
+    return stage_onehot
+
+
+class SARMRewardModel(PreTrainedRewardModel):
+    """
+    SARM Reward Model for stage-aware task completion rewards.
+
+    Uses two separate transformer models:
+    - StageTransformer: Classifies which stage/subtask
+    - SubtaskTransformer: Predicts within-stage progress (tau)
+
+    Training uses 75%/25% GT/predicted stage conditioning (teacher forcing).
+    """
+
+    name = "sarm"
+    config_class = SARMConfig
+
+    def __init__(self, config: SARMConfig, dataset_stats: dict | None = None, dataset_meta=None, **kwargs):
+        super().__init__(config)
+        config.validate_features()
+        self.config = config
+        self.dataset_stats = dataset_stats
+        self.device = torch.device(
+            config.device if config.device else "cuda" if torch.cuda.is_available() else "cpu"
+        )
+
+        # Load temporal proportions based on annotation_mode
+        if config.annotation_mode == "single_stage":
+            logging.info(f"Using single_stage mode: sparse_subtask_names={config.sparse_subtask_names}")
+        elif dataset_meta is not None:
+            self._load_temporal_proportions(dataset_meta)
+
+        # Create two separate models
+        self.stage_model = StageTransformer(
+            d_model=config.hidden_dim,
+            vis_emb_dim=config.image_dim,
+            text_emb_dim=config.text_dim,
+            state_dim=config.max_state_dim,
+            n_layers=config.num_layers,
+            n_heads=config.num_heads,
+            dropout=config.dropout,
+            num_cameras=1,
+            num_classes_sparse=config.num_sparse_stages,
+            num_classes_dense=config.num_dense_stages or config.num_sparse_stages,
+        )
+
+        self.subtask_model = SubtaskTransformer(
+            d_model=config.hidden_dim,
+            vis_emb_dim=config.image_dim,
+            text_emb_dim=config.text_dim,
+            state_dim=config.max_state_dim,
+            n_layers=config.num_layers,
+            n_heads=config.num_heads,
+            dropout=config.dropout,
+            num_cameras=1,
+        )
+
+        self.stage_model.to(self.device)
+        self.subtask_model.to(self.device)
+
+        self.gt_stage_ratio = 0.75
+
+        if config.uses_dual_heads:
+            logging.info(
+                f"SARM initialized with dual heads: {config.num_sparse_stages} sparse stages, "
+                f"{config.num_dense_stages} dense stages"
+            )
+        else:
+            logging.info(f"SARM initialized with sparse head only: {config.num_sparse_stages} stages")
+
+        logging.info(f"SARM initialized on {self.device}")
+
+    def _load_proportions_from_json(self, path, annotation_type: str) -> tuple[list[str], list[float]]:
+        """Load temporal proportions from a JSON file (preserving order)."""
+        if not path.exists():
+            raise ValueError(
+                f"{annotation_type.capitalize()} temporal proportions not found at {path}. "
+                f"Run the subtask annotation tool with --{annotation_type}-subtasks to generate annotations."
+            )
+        with open(path) as f:
+            proportions_dict = json.load(f)
+        names = list(proportions_dict.keys())
+        logging.info(f"Loaded {len(names)} {annotation_type} subtasks: {names}")
+        logging.info(f"{annotation_type.capitalize()} temporal proportions: {proportions_dict}")
+        return names, [proportions_dict[name] for name in names]
+
+    def _load_temporal_proportions(self, dataset_meta) -> None:
+        """Load temporal proportions based on annotation_mode."""
+        meta_path = dataset_meta.root / "meta"
+
+        if self.config.annotation_mode == "dual":
+            names, props = self._load_proportions_from_json(
+                meta_path / "temporal_proportions_sparse.json", "sparse"
+            )
+            (
+                self.config.num_sparse_stages,
+                self.config.sparse_subtask_names,
+                self.config.sparse_temporal_proportions,
+            ) = len(names), names, props
+
+        if self.config.annotation_mode in ["dense_only", "dual"]:
+            names, props = self._load_proportions_from_json(
+                meta_path / "temporal_proportions_dense.json", "dense"
+            )
+            (
+                self.config.num_dense_stages,
+                self.config.dense_subtask_names,
+                self.config.dense_temporal_proportions,
+            ) = len(names), names, props
+            if self.config.annotation_mode == "dense_only":
+                logging.info(f"Using auto-generated sparse 'task' stage: {self.config.sparse_subtask_names}")
+
+    def to(self, device):
+        """Override to method to ensure all components move together."""
+        super().to(device)
+        self.device = device if isinstance(device, torch.device) else torch.device(device)
+        self.stage_model.to(device)
+        self.subtask_model.to(device)
+        return self
+
+    def compute_reward(self, batch: dict[str, Tensor]) -> Tensor:
+        """Compute dense progress reward in [0, 1] from batch.
+
+        Expects batch to contain:
+        - "observation_features" or video embeddings: (B, T, 512)
+        - "language_embedding" or text embeddings: (B, 512)
+        - optionally "observation.state": (B, T, state_dim)
+        """
+        text_emb = batch.get("language_embedding", batch.get("text_features"))
+        video_emb = batch.get("observation_features", batch.get("video_features"))
+        state = batch.get("observation.state", batch.get("state_features"))
+
+        rewards = self.calculate_rewards(text_emb, video_emb, state)
+        if isinstance(rewards, np.ndarray):
+            rewards = torch.from_numpy(rewards).float()
+        return rewards
+
+    @torch.no_grad()
+    def calculate_rewards(
+        self,
+        text_embeddings: np.ndarray | torch.Tensor,
+        video_embeddings: np.ndarray | torch.Tensor,
+        state_features: np.ndarray | torch.Tensor | None = None,
+        lengths: np.ndarray | torch.Tensor | None = None,
+        return_all_frames: bool = False,
+        return_stages: bool = False,
+        return_confidence: bool = False,
+        head_mode: str | None = "sparse",
+        frame_index: int | None = None,
+    ) -> np.ndarray | tuple:
+        """
+        Calculate rewards for given text, video, and state representations.
+
+        This is the canonical method for SARM reward computation, used for:
+        - Inference/visualization
+        - RA-BC weight computation
+        """
+        if isinstance(text_embeddings, np.ndarray):
+            text_embeddings = torch.tensor(text_embeddings, dtype=torch.float32)
+        if isinstance(video_embeddings, np.ndarray):
+            video_embeddings = torch.tensor(video_embeddings, dtype=torch.float32)
+        if state_features is not None and isinstance(state_features, np.ndarray):
+            state_features = torch.tensor(state_features, dtype=torch.float32)
+
+        if text_embeddings.dim() == 1:
+            text_embeddings = text_embeddings.unsqueeze(0)
+            video_embeddings = video_embeddings.unsqueeze(0)
+            if state_features is not None:
+                state_features = state_features.unsqueeze(0)
+            single_sample = True
+        else:
+            single_sample = False
+
+        batch_size = video_embeddings.shape[0]
+        seq_len = video_embeddings.shape[1]
+
+        scheme = head_mode
+
+        if lengths is None:
+            lengths = torch.full((batch_size,), seq_len, dtype=torch.int32)
+        elif isinstance(lengths, np.ndarray):
+            lengths = torch.tensor(lengths, dtype=torch.int32)
+
+        img_seq = video_embeddings.unsqueeze(1).to(self.device)
+        lang_emb = text_embeddings.to(self.device)
+        state = (
+            state_features.to(self.device)
+            if state_features is not None
+            else torch.zeros(batch_size, seq_len, self.config.max_state_dim, device=self.device)
+        )
+        lens = lengths.to(self.device)
+
+        state = pad_state_to_max_dim(state, self.config.max_state_dim)
+
+        num_classes = self.config.num_sparse_stages if scheme == "sparse" else self.config.num_dense_stages
+
+        stage_logits = self.stage_model(img_seq, lang_emb, state, lens, scheme=scheme)
+        stage_probs = F.softmax(stage_logits, dim=-1)
+        stage_idx = stage_probs.argmax(dim=-1)
+        stage_conf = stage_probs.gather(-1, stage_idx.unsqueeze(-1)).squeeze(-1)
+
+        stage_onehot = F.one_hot(stage_idx, num_classes=num_classes).float()
+        stage_emb = stage_onehot.unsqueeze(1)
+
+        tau_pred = self.subtask_model(img_seq, lang_emb, state, lens, stage_emb, scheme=scheme)
+
+        raw_reward = stage_idx.float() + tau_pred
+
+        if scheme == "sparse":
+            normalized_reward = normalize_stage_tau(
+                raw_reward,
+                num_stages=num_classes,
+                temporal_proportions=self.config.sparse_temporal_proportions,
+                subtask_names=self.config.sparse_subtask_names,
+            )
+        else:
+            normalized_reward = normalize_stage_tau(
+                raw_reward,
+                num_stages=num_classes,
+                temporal_proportions=self.config.dense_temporal_proportions,
+                subtask_names=self.config.dense_subtask_names,
+            )
+
+        if frame_index is None:
+            frame_index = self.config.n_obs_steps
+
+        if return_all_frames:
+            rewards = normalized_reward.cpu().numpy()
+        else:
+            rewards = normalized_reward[:, frame_index].cpu().numpy()
+
+        if single_sample:
+            rewards = rewards[0] if not return_all_frames else rewards[0]
+
+        outputs = [rewards]
+        if return_stages:
+            probs = stage_probs.cpu().numpy()
+            if single_sample:
+                probs = probs[0]
+            outputs.append(probs)
+        if return_confidence:
+            conf = stage_conf.cpu().numpy()
+            if single_sample:
+                conf = conf[0]
+            outputs.append(conf)
+
+        return outputs[0] if len(outputs) == 1 else tuple(outputs)
+
+    def train(self, mode: bool = True):
+        """Set training mode for both models."""
+        super().train(mode)
+        self.stage_model.train(mode)
+        self.subtask_model.train(mode)
+        return self
+
+    def eval(self):
+        """Set evaluation mode for both models."""
+        return self.train(False)
+
+    def parameters(self):
+        """Override to return trainable parameters from both models."""
+        from itertools import chain
+
+        return chain(self.stage_model.parameters(), self.subtask_model.parameters())
+
+    def get_optim_params(self):
+        """Override to return optimizer parameters from both models."""
+        return self.parameters()
+
+    def reset(self):
+        pass
+
+    def _train_step(
+        self,
+        img_emb: torch.Tensor,
+        lang_emb: torch.Tensor,
+        state: torch.Tensor,
+        lengths: torch.Tensor,
+        targets: torch.Tensor,
+        scheme: str,
+    ) -> dict[str, torch.Tensor]:
+        """Single training step for one annotation scheme."""
+        num_classes = self.config.num_sparse_stages if scheme == "sparse" else self.config.num_dense_stages
+
+        gt_stage = torch.floor(targets).long().clamp(0, num_classes - 1)
+        gt_tau = torch.remainder(targets, 1.0)
+
+        stage_pred = self.stage_model(img_emb, lang_emb, state, lengths, scheme=scheme)
+
+        if random.random() < self.gt_stage_ratio:
+            stage_emb = gen_stage_emb(num_classes, targets)
+        else:
+            stage_idx = stage_pred.argmax(dim=-1)
+            stage_onehot = F.one_hot(stage_idx, num_classes=num_classes).float()
+            stage_emb = stage_onehot.unsqueeze(1)
+
+        tau_pred = self.subtask_model(img_emb, lang_emb, state, lengths, stage_emb, scheme=scheme)
+
+        stage_loss = F.cross_entropy(stage_pred.view(-1, num_classes), gt_stage.view(-1), reduction="mean")
+        subtask_loss = F.mse_loss(tau_pred, gt_tau, reduction="mean")
+
+        return {
+            "stage_loss": stage_loss,
+            "subtask_loss": subtask_loss,
+            "total_loss": stage_loss + subtask_loss,
+        }
+
+    def forward(self, batch):
+        """Forward pass for SARM reward model training."""
+        observation = batch.get(OBS_STR, batch)
+
+        video_features = observation["video_features"].to(self.device)
+        text_features = observation["text_features"].to(self.device)
+        state_features = observation.get("state_features")
+        if state_features is not None:
+            state_features = state_features.to(self.device)
+
+        batch_size = video_features.shape[0]
+        seq_len = video_features.shape[1]
+
+        lengths = observation.get("lengths")
+        if lengths is None:
+            lengths = torch.full((batch_size,), seq_len, dtype=torch.int32, device=self.device)
+        else:
+            lengths = lengths.to(self.device)
+
+        img_emb = video_features.unsqueeze(1)
+
+        if state_features is None:
+            state_features = torch.zeros(batch_size, seq_len, self.config.max_state_dim, device=self.device)
+        else:
+            state_features = pad_state_to_max_dim(state_features, self.config.max_state_dim)
+
+        output_dict = {}
+        total_loss = torch.tensor(0.0, device=self.device)
+
+        sparse_targets = observation.get("sparse_targets")
+        if sparse_targets is None:
+            sparse_targets = observation.get("targets")
+        if sparse_targets is None:
+            raise ValueError("sparse_targets (or targets) is required for SARM training")
+        sparse_targets = sparse_targets.to(self.device)
+
+        sparse_result = self._train_step(
+            img_emb, text_features, state_features, lengths, sparse_targets, scheme="sparse"
+        )
+        output_dict["sparse_stage_loss"] = sparse_result["stage_loss"].item()
+        output_dict["sparse_subtask_loss"] = sparse_result["subtask_loss"].item()
+        total_loss = total_loss + sparse_result["total_loss"]
+
+        if self.config.uses_dual_heads:
+            dense_targets = observation.get("dense_targets")
+            if dense_targets is not None:
+                dense_targets = dense_targets.to(self.device)
+                dense_result = self._train_step(
+                    img_emb, text_features, state_features, lengths, dense_targets, scheme="dense"
+                )
+                output_dict["dense_stage_loss"] = dense_result["stage_loss"].item()
+                output_dict["dense_subtask_loss"] = dense_result["subtask_loss"].item()
+                total_loss = total_loss + dense_result["total_loss"]
+
+        output_dict["total_loss"] = total_loss.item()
+        return total_loss, output_dict
+
+
+def compute_stage_loss(stage_logits: torch.Tensor, target_stages: torch.Tensor) -> torch.Tensor:
+    """Compute cross-entropy loss for stage classification."""
+    _, _, num_stages = stage_logits.shape
+    stage_logits_flat = stage_logits.reshape(-1, num_stages)
+    target_stages_flat = target_stages.reshape(-1).clamp(0, num_stages - 1)
+    return F.cross_entropy(stage_logits_flat, target_stages_flat)

src/lerobot/rewards/sarm/processor_sarm.py

@@ -0,0 +1,463 @@
+#!/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.
+
+"""SARM Processor for encoding images/text and generating stage+tau targets."""
+
+import random
+from typing import Any
+
+import numpy as np
+import pandas as pd
+import torch
+from faker import Faker
+from PIL import Image
+from transformers import CLIPModel, CLIPProcessor
+
+from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    RenameObservationsProcessorStep,
+)
+from lerobot.processor.converters import (
+    from_tensor_to_numpy,
+    policy_action_to_transition,
+    transition_to_policy_action,
+)
+from lerobot.processor.core import EnvTransition, TransitionKey
+from lerobot.processor.pipeline import PipelineFeatureType
+from lerobot.rewards.sarm.configuration_sarm import SARMConfig
+from lerobot.rewards.sarm.sarm_utils import (
+    apply_rewind_augmentation,
+    compute_absolute_indices,
+    find_stage_and_tau,
+    pad_state_to_max_dim,
+)
+from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+
+
+class SARMEncodingProcessorStep(ProcessorStep):
+    """ProcessorStep that encodes images and text with CLIP and generates stage and progress labels for SARM."""
+
+    def __init__(
+        self,
+        config: SARMConfig,
+        image_key: str | None = None,
+        dataset_meta=None,
+        dataset_stats: dict | None = None,
+    ):
+        super().__init__()
+        self.config = config
+        self.image_key = image_key or config.image_key
+        self.dataset_meta = dataset_meta
+        self.dataset_stats = dataset_stats
+        self.annotation_mode = config.annotation_mode
+
+        def make_props_dict(names, props):
+            return dict(zip(names, props, strict=True)) if names and props else None
+
+        self.sparse_temporal_proportions = make_props_dict(
+            config.sparse_subtask_names, config.sparse_temporal_proportions
+        )
+        self.sparse_subtask_names = config.sparse_subtask_names
+
+        self.dense_subtask_names = config.dense_subtask_names if config.uses_dual_heads else None
+        self.dense_temporal_proportions = (
+            make_props_dict(config.dense_subtask_names, config.dense_temporal_proportions)
+            if config.uses_dual_heads
+            else None
+        )
+
+        self.device = torch.device(
+            self.config.device if self.config.device else "cuda" if torch.cuda.is_available() else "cpu"
+        )
+
+        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", use_fast=True)
+        self.clip_model.to(self.device)
+        self.clip_model.eval()
+
+        self.verbs = ["move", "grasp", "rotate", "push", "pull", "slide", "lift", "place"]
+        self.fake = Faker()
+
+    def _find_episode_for_frame(self, frame_idx: int) -> int:
+        """Find the episode index for a given frame index."""
+        for ep_idx in range(len(self.dataset_meta.episodes)):
+            ep_start = self.dataset_meta.episodes[ep_idx]["dataset_from_index"]
+            ep_end = self.dataset_meta.episodes[ep_idx]["dataset_to_index"]
+            if ep_start <= frame_idx < ep_end:
+                return ep_idx
+        return 0
+
+    def _get_episode_indices(self, frame_indices: np.ndarray, episode_index) -> np.ndarray:
+        """Get episode indices for each frame index."""
+        if episode_index is None:
+            return np.array([self._find_episode_for_frame(int(f)) for f in frame_indices])
+
+        episode_indices = np.atleast_1d(np.asarray(from_tensor_to_numpy(episode_index)))
+
+        if len(episode_indices) == 1 and len(frame_indices) > 1:
+            return np.array([self._find_episode_for_frame(int(f)) for f in frame_indices])
+
+        return episode_indices
+
+    def _generate_perturbed_task(self) -> str:
+        """Generate a random perturbed task string for language perturbation."""
+        num_words = random.randint(1, 5)
+        verb = random.choice(self.verbs)
+        phrase = " ".join([verb] + self.fake.words(nb=num_words))
+        return phrase
+
+    def _get_annotation_config(self, annotation_type: str) -> tuple[list[str], dict[str, float] | None]:
+        """Get global subtask names and temporal proportions for an annotation type."""
+        if annotation_type == "dense":
+            return self.dense_subtask_names, self.dense_temporal_proportions
+        return self.sparse_subtask_names, self.sparse_temporal_proportions
+
+    def _load_episode_annotations(
+        self,
+        ep_idx: int,
+        episodes_df: pd.DataFrame | None,
+        annotation_type: str,
+        global_names: list[str],
+    ) -> tuple[list | None, list | None, list | None]:
+        """Load subtask annotations for an episode from DataFrame."""
+        if episodes_df is None or len(global_names) == 1:
+            return None, None, None
+
+        def col(suffix):
+            prefixed = f"{annotation_type}_{suffix}"
+            return prefixed if prefixed in episodes_df.columns else suffix
+
+        col_names = col("subtask_names")
+        if col_names not in episodes_df.columns or ep_idx >= len(episodes_df):
+            return None, None, None
+
+        subtask_names = episodes_df.loc[ep_idx, col_names]
+        if subtask_names is None or (isinstance(subtask_names, float) and pd.isna(subtask_names)):
+            return None, None, None
+
+        return (
+            subtask_names,
+            episodes_df.loc[ep_idx, col("subtask_start_frames")],
+            episodes_df.loc[ep_idx, col("subtask_end_frames")],
+        )
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        """Encode images, text, and normalize states in the transition."""
+        new_transition = transition.copy() if hasattr(transition, "copy") else dict(transition)
+        observation = new_transition.get(TransitionKey.OBSERVATION)
+        comp_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
+
+        frame_index = comp_data.get("index")
+        episode_index = comp_data.get("episode_index")
+
+        if frame_index is None:
+            raise ValueError("Frame index ('index') not found in COMPLEMENTARY_DATA")
+        if episode_index is None:
+            raise ValueError("Episode index ('episode_index') not found in COMPLEMENTARY_DATA")
+
+        frame_indices = np.atleast_1d(np.asarray(from_tensor_to_numpy(frame_index)))
+        episode_indices = self._get_episode_indices(frame_indices, episode_index)
+
+        image = observation.get(self.image_key)
+        if isinstance(image, torch.Tensor):
+            image = image.cpu().numpy()
+
+        if image.ndim == 4:
+            image = image[np.newaxis, ...]
+        elif image.ndim == 3:
+            image = image[np.newaxis, np.newaxis, ...]
+
+        batch_size = image.shape[0]
+        total_frames = image.shape[1]
+        n_obs_steps = self.config.n_obs_steps
+        max_rewind_steps = self.config.max_rewind_steps
+        n_obs_frames = 1 + n_obs_steps
+
+        rewind_steps = torch.zeros(batch_size, dtype=torch.int32)
+        apply_rewind = self.training and random.random() < self.config.rewind_probability
+
+        if apply_rewind and self.dataset_meta is not None:
+            for b_idx, (ep_idx, frame_idx) in enumerate(
+                zip(episode_indices.tolist(), frame_indices.tolist(), strict=True)
+            ):
+                ep_idx, frame_idx = int(ep_idx), int(frame_idx)
+                ep_start = self.dataset_meta.episodes[ep_idx]["dataset_from_index"]
+
+                rewind_step, _ = apply_rewind_augmentation(
+                    frame_idx, ep_start, n_obs_steps, max_rewind_steps, frame_gap=self.config.frame_gap
+                )
+                rewind_steps[b_idx] = rewind_step
+
+        lengths = n_obs_frames + rewind_steps
+
+        for b_idx in range(batch_size):
+            valid_len = lengths[b_idx].item()
+            if valid_len < total_frames:
+                image[b_idx, valid_len:] = 0
+
+        video_features = self._encode_images_batch(image)
+        observation["video_features"] = video_features
+
+        state_key = self.config.state_key
+        state_data = observation.get(state_key)
+
+        if isinstance(state_data, torch.Tensor):
+            state_tensor = state_data.float()
+        else:
+            state_tensor = torch.tensor(state_data, dtype=torch.float32)
+
+        if state_tensor.ndim == 2:
+            state_tensor = state_tensor.unsqueeze(0)
+        elif state_tensor.ndim == 1:
+            state_tensor = state_tensor.unsqueeze(0).unsqueeze(0)
+
+        for b_idx in range(batch_size):
+            valid_len = lengths[b_idx].item()
+            if valid_len < state_tensor.shape[1]:
+                state_tensor[b_idx, valid_len:] = 0
+
+        observation["state_features"] = pad_state_to_max_dim(state_tensor, self.config.max_state_dim)
+
+        task = comp_data.get("task")
+        if isinstance(task, list):
+            task = task[0] if task else ""
+
+        apply_perturbation = self.training and random.random() < self.config.language_perturbation_probability
+        if apply_perturbation:
+            task = self._generate_perturbed_task()
+
+        observation["text_features"] = self._encode_text_clip(task, batch_size)
+
+        observation["lengths"] = lengths
+
+        if self.dataset_meta is not None:
+            episodes_df = self.dataset_meta.episodes.to_pandas()
+
+            if self.sparse_temporal_proportions is not None:
+                if apply_perturbation:
+                    sparse_targets = torch.zeros(batch_size, total_frames, dtype=torch.float32)
+                else:
+                    sparse_targets = self._compute_batch_targets(
+                        frame_indices, episode_indices, lengths, rewind_steps, episodes_df, "sparse"
+                    )
+                observation["sparse_targets"] = sparse_targets
+
+            if self.config.uses_dual_heads and self.dense_temporal_proportions is not None:
+                if apply_perturbation:
+                    dense_targets = torch.zeros(batch_size, total_frames, dtype=torch.float32)
+                else:
+                    dense_targets = self._compute_batch_targets(
+                        frame_indices, episode_indices, lengths, rewind_steps, episodes_df, "dense"
+                    )
+                observation["dense_targets"] = dense_targets
+
+        new_transition[TransitionKey.OBSERVATION] = observation
+        return new_transition
+
+    def _compute_batch_targets(
+        self,
+        frame_indices: np.ndarray,
+        episode_indices: np.ndarray,
+        lengths: torch.Tensor,
+        rewind_steps: torch.Tensor,
+        episodes_df: pd.DataFrame | None,
+        annotation_type: str,
+    ) -> torch.Tensor:
+        """Compute stage+tau targets for a batch of samples."""
+        batch_size = len(frame_indices)
+        n_obs_steps = self.config.n_obs_steps
+        max_rewind_steps = self.config.max_rewind_steps
+        total_frames = 1 + n_obs_steps + max_rewind_steps
+        frame_gap = self.config.frame_gap
+
+        global_names, temporal_props = self._get_annotation_config(annotation_type)
+        targets = torch.zeros(batch_size, total_frames, dtype=torch.float32)
+
+        for b_idx in range(batch_size):
+            ep_idx = int(episode_indices[b_idx])
+            frame_idx = int(frame_indices[b_idx])
+
+            ep_start = self.dataset_meta.episodes[ep_idx]["dataset_from_index"]
+            ep_end = self.dataset_meta.episodes[ep_idx]["dataset_to_index"]
+            ep_length = ep_end - ep_start
+
+            subtask_names, subtask_start_frames, subtask_end_frames = self._load_episode_annotations(
+                ep_idx, episodes_df, annotation_type, global_names
+            )
+
+            obs_indices, _ = compute_absolute_indices(
+                frame_idx, ep_start, ep_end, n_obs_steps, frame_gap=frame_gap
+            )
+            obs_indices = obs_indices.tolist()
+
+            for t_idx, abs_idx in enumerate(obs_indices):
+                rel_frame = abs_idx - ep_start
+                targets[b_idx, t_idx] = find_stage_and_tau(
+                    rel_frame,
+                    ep_length,
+                    subtask_names,
+                    subtask_start_frames,
+                    subtask_end_frames,
+                    global_names,
+                    temporal_props,
+                    return_combined=True,
+                )
+
+            rewind_step = rewind_steps[b_idx].item()
+            if rewind_step > 0:
+                _, rewind_indices = apply_rewind_augmentation(
+                    frame_idx,
+                    ep_start,
+                    n_obs_steps,
+                    max_rewind_steps,
+                    frame_gap=frame_gap,
+                    rewind_step=rewind_step,
+                )
+
+                for r_idx, abs_idx in enumerate(rewind_indices[:rewind_step]):
+                    rel_frame = max(0, abs_idx - ep_start)
+                    targets[b_idx, n_obs_steps + 1 + r_idx] = find_stage_and_tau(
+                        rel_frame,
+                        ep_length,
+                        subtask_names,
+                        subtask_start_frames,
+                        subtask_end_frames,
+                        global_names,
+                        temporal_props,
+                        return_combined=True,
+                    )
+
+        return targets
+
+    @property
+    def training(self) -> bool:
+        return getattr(self, "_training_mode", True)
+
+    def train(self, mode: bool = True):
+        """Set training mode for augmentation decisions."""
+        self._training_mode = mode
+        return self
+
+    def eval(self):
+        """Set evaluation mode (disable augmentations)."""
+        return self.train(False)
+
+    @torch.no_grad()
+    def _encode_images_batch(self, images: np.ndarray) -> torch.Tensor:
+        """Encode a batch of images using CLIP."""
+        batch_size, seq_length = images.shape[0], images.shape[1]
+        images = images.reshape(batch_size * seq_length, *images.shape[2:])
+
+        num_frames = images.shape[0]
+        images_list = []
+        for i in range(num_frames):
+            img = images[i]
+            if img.shape[0] in [1, 3]:
+                img = img.transpose(1, 2, 0)
+
+            if img.shape[-1] == 1:
+                img = np.repeat(img, 3, axis=-1)
+
+            if img.dtype != np.uint8:
+                img = (img * 255).astype(np.uint8) if img.max() <= 1.0 else img.astype(np.uint8)
+
+            images_list.append(Image.fromarray(img))
+
+        all_embeddings = []
+        for i in range(0, num_frames, self.config.clip_batch_size):
+            batch_imgs = images_list[i : i + self.config.clip_batch_size]
+
+            inputs = self.clip_processor(images=batch_imgs, return_tensors="pt")
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+            embeddings = self.clip_model.get_image_features(**inputs).detach().cpu()
+
+            if embeddings.dim() == 1:
+                embeddings = embeddings.unsqueeze(0)
+
+            all_embeddings.append(embeddings)
+
+        all_embeddings = torch.cat(all_embeddings)
+        all_embeddings = all_embeddings.reshape(batch_size, seq_length, -1)
+
+        return all_embeddings
+
+    @torch.no_grad()
+    def _encode_text_clip(self, text: str, batch_size: int) -> torch.Tensor:
+        """Encode text using CLIP text encoder (per SARM paper A.4)."""
+        inputs = self.clip_processor.tokenizer([text], return_tensors="pt", padding=True, truncation=True)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        text_embedding = self.clip_model.get_text_features(**inputs).detach().cpu()
+        text_embedding = text_embedding.expand(batch_size, -1)
+
+        return text_embedding
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        """Add encoded features to the observation features."""
+        features[PipelineFeatureType.OBSERVATION]["video_features"] = PolicyFeature(
+            type=FeatureType.VISUAL, shape=(self.config.num_frames, self.config.image_dim)
+        )
+        features[PipelineFeatureType.OBSERVATION]["text_features"] = PolicyFeature(
+            type=FeatureType.LANGUAGE, shape=(self.config.text_dim,)
+        )
+        features[PipelineFeatureType.OBSERVATION]["state_features"] = PolicyFeature(
+            type=FeatureType.STATE, shape=(self.config.num_frames, self.config.max_state_dim)
+        )
+        return features
+
+
+def make_sarm_pre_post_processors(
+    config: SARMConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+    dataset_meta=None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """Create pre-processor and post-processor pipelines for SARM."""
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=[
+                AddBatchDimensionProcessorStep(),
+                RenameObservationsProcessorStep(rename_map={}),
+                NormalizerProcessorStep(
+                    features={**config.input_features, **config.output_features},
+                    norm_map=config.normalization_mapping,
+                    stats=dataset_stats,
+                ),
+                SARMEncodingProcessorStep(
+                    config=config, dataset_meta=dataset_meta, dataset_stats=dataset_stats
+                ),
+                DeviceProcessorStep(device=config.device),
+            ],
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=[DeviceProcessorStep(device="cpu")],
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/rewards/sarm/sarm_utils.py

@@ -0,0 +1,230 @@
+#!/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.
+
+import random
+
+import numpy as np
+import torch
+import torch.nn.functional as F  # noqa: N812
+
+
+def find_stage_and_tau(
+    current_frame: int,
+    episode_length: int,
+    subtask_names: list | None,
+    subtask_start_frames: list | None,
+    subtask_end_frames: list | None,
+    global_subtask_names: list,
+    temporal_proportions: dict,
+    return_combined: bool = False,
+) -> tuple[int, float] | float:
+    """Find stage and within-stage progress (tau) for a frame.
+
+    Args:
+        current_frame: Frame index relative to episode start
+        episode_length: Total frames in episode
+        subtask_names: Subtask names for this episode (None for single_stage)
+        subtask_start_frames: Subtask start frames
+        subtask_end_frames: Subtask end frames
+        global_subtask_names: Global list of all subtask names
+        temporal_proportions: Dict of temporal proportions
+        return_combined: If True, return stage+tau as float; else (stage_idx, tau) tuple
+
+    Returns:
+        Float (stage.tau) if return_combined, else (stage_idx, tau) tuple
+    """
+    stage_idx, tau = 0, 0.0
+    num_stages = len(global_subtask_names)
+
+    # Single-stage mode: linear progress from 0 to 1
+    if num_stages == 1:
+        tau = min(1.0, max(0.0, current_frame / max(episode_length - 1, 1)))
+    elif subtask_names is None:
+        pass  # stage_idx=0, tau=0.0
+    elif current_frame < subtask_start_frames[0]:
+        pass  # Before first subtask: stage_idx=0, tau=0.0
+    elif current_frame > subtask_end_frames[-1]:
+        stage_idx, tau = num_stages - 1, 0.999  # After last subtask
+    else:
+        # Find which subtask this frame belongs to
+        found = False
+        for name, start, end in zip(subtask_names, subtask_start_frames, subtask_end_frames, strict=True):
+            if start <= current_frame <= end:
+                stage_idx = global_subtask_names.index(name) if name in global_subtask_names else 0
+                tau = compute_tau(current_frame, start, end)
+                found = True
+                break
+        # Frame between subtasks - use previous subtask's end state
+        if not found:
+            for j in range(len(subtask_names) - 1):
+                if subtask_end_frames[j] < current_frame < subtask_start_frames[j + 1]:
+                    name = subtask_names[j]
+                    stage_idx = global_subtask_names.index(name) if name in global_subtask_names else j
+                    tau = 1.0
+                    break
+
+    if return_combined:
+        # Clamp to avoid overflow at end
+        if stage_idx >= num_stages - 1 and tau >= 1.0:
+            return num_stages - 1 + 0.999
+        return stage_idx + tau
+    return stage_idx, tau
+
+
+def compute_absolute_indices(
+    frame_idx: int,
+    ep_start: int,
+    ep_end: int,
+    n_obs_steps: int,
+    frame_gap: int = 30,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute absolute frame indices with clamping for bidirectional observation sequence."""
+    half_steps = n_obs_steps // 2
+
+    past_deltas = [-frame_gap * i for i in range(half_steps, 0, -1)]
+    future_deltas = [frame_gap * i for i in range(1, half_steps + 1)]
+    delta_indices = past_deltas + [0] + future_deltas
+
+    frames = []
+    out_of_bounds = []
+
+    for delta in delta_indices:
+        target_idx = frame_idx + delta
+        clamped_idx = max(ep_start, min(ep_end - 1, target_idx))
+        frames.append(clamped_idx)
+        out_of_bounds.append(1 if target_idx != clamped_idx else 0)
+
+    return torch.tensor(frames), torch.tensor(out_of_bounds)
+
+
+def apply_rewind_augmentation(
+    frame_idx: int,
+    ep_start: int,
+    n_obs_steps: int,
+    max_rewind_steps: int,
+    frame_gap: int = 30,
+    rewind_step: int | None = None,
+) -> tuple[int, list[int]]:
+    """Generate rewind frame indices for temporal augmentation."""
+    half_steps = n_obs_steps // 2
+    earliest_obs_frame = frame_idx - half_steps * frame_gap
+
+    if earliest_obs_frame <= ep_start:
+        return 0, []
+
+    available_history = earliest_obs_frame - ep_start
+    max_valid_step = available_history // frame_gap
+    max_rewind = min(max_rewind_steps, max(0, max_valid_step))
+
+    if max_rewind <= 0:
+        return 0, []
+
+    rewind_step = random.randint(1, max_rewind) if rewind_step is None else min(rewind_step, max_rewind)
+
+    if rewind_step == 0:
+        return 0, []
+
+    rewind_indices = []
+    for i in range(1, rewind_step + 1):
+        idx = earliest_obs_frame - i * frame_gap
+        idx = max(ep_start, idx)
+        rewind_indices.append(idx)
+
+    return rewind_step, rewind_indices
+
+
+def compute_tau(current_frame: int | float, subtask_start: int | float, subtask_end: int | float) -> float:
+    """Compute τ_t = (t - s_k) / (e_k - s_k) ∈ [0, 1]. Returns 1.0 for zero-duration subtasks."""
+    duration = subtask_end - subtask_start
+    if duration <= 0:
+        return 1.0
+    return float(np.clip((current_frame - subtask_start) / duration, 0.0, 1.0))
+
+
+def pad_state_to_max_dim(state: torch.Tensor, max_state_dim: int) -> torch.Tensor:
+    """Pad the state tensor's last dimension to max_state_dim with zeros."""
+    current_dim = state.shape[-1]
+    if current_dim >= max_state_dim:
+        return state[..., :max_state_dim]
+
+    padding = (0, max_state_dim - current_dim)
+    return F.pad(state, padding, mode="constant", value=0)
+
+
+def temporal_proportions_to_breakpoints(
+    temporal_proportions: dict[str, float] | list[float] | None,
+    subtask_names: list[str] | None = None,
+) -> list[float] | None:
+    """Convert temporal proportions to cumulative breakpoints for normalization."""
+    if temporal_proportions is None:
+        return None
+
+    if isinstance(temporal_proportions, dict):
+        if subtask_names is not None:
+            proportions = [temporal_proportions.get(name, 0.0) for name in subtask_names]
+        else:
+            proportions = list(temporal_proportions.values())
+    else:
+        proportions = list(temporal_proportions)
+
+    total = sum(proportions)
+    if total > 0 and abs(total - 1.0) > 1e-6:
+        proportions = [p / total for p in proportions]
+
+    breakpoints = [0.0]
+    cumsum = 0.0
+    for prop in proportions:
+        cumsum += prop
+        breakpoints.append(cumsum)
+    breakpoints[-1] = 1.0
+
+    return breakpoints
+
+
+def normalize_stage_tau(
+    x: float | torch.Tensor,
+    num_stages: int | None = None,
+    breakpoints: list[float] | None = None,
+    temporal_proportions: dict[str, float] | list[float] | None = None,
+    subtask_names: list[str] | None = None,
+) -> float | torch.Tensor:
+    """Normalize stage+tau reward to [0, 1] with custom breakpoints."""
+    if breakpoints is not None:
+        num_stages = len(breakpoints) - 1
+    elif temporal_proportions is not None:
+        breakpoints = temporal_proportions_to_breakpoints(temporal_proportions, subtask_names)
+        num_stages = len(breakpoints) - 1
+    elif num_stages is not None:
+        breakpoints = [i / num_stages for i in range(num_stages + 1)]
+    else:
+        raise ValueError("Either num_stages, breakpoints, or temporal_proportions must be provided")
+
+    if isinstance(x, torch.Tensor):
+        result = torch.zeros_like(x)
+        for i in range(num_stages):
+            mask = (x >= i) & (x < i + 1)
+            tau_in_stage = x - i
+            result[mask] = breakpoints[i] + tau_in_stage[mask] * (breakpoints[i + 1] - breakpoints[i])
+        result[x >= num_stages] = 1.0
+        return result.clamp(0.0, 1.0)
+    else:
+        if x < 0:
+            return 0.0
+        if x >= num_stages:
+            return 1.0
+        stage = int(x)
+        tau = x - stage
+        return breakpoints[stage] + tau * (breakpoints[stage + 1] - breakpoints[stage])