From 0508745852ff265b1b8cfd66f9126a58e0a99dda Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Thu, 12 Mar 2026 15:19:10 +0100 Subject: [PATCH] refactor(rewards): migrate SARM from policies/sarm/ to rewards/sarm/ --- src/lerobot/rewards/sarm/__init__.py | 0 .../rewards/sarm/compute_rabc_weights.py | 870 ++++++++++++++++++ .../rewards/sarm/configuration_sarm.py | 232 +++++ src/lerobot/rewards/sarm/modeling_sarm.py | 627 +++++++++++++ src/lerobot/rewards/sarm/processor_sarm.py | 463 ++++++++++ src/lerobot/rewards/sarm/sarm_utils.py | 230 +++++ 6 files changed, 2422 insertions(+) create mode 100644 src/lerobot/rewards/sarm/__init__.py create mode 100644 src/lerobot/rewards/sarm/compute_rabc_weights.py create mode 100644 src/lerobot/rewards/sarm/configuration_sarm.py create mode 100644 src/lerobot/rewards/sarm/modeling_sarm.py create mode 100644 src/lerobot/rewards/sarm/processor_sarm.py create mode 100644 src/lerobot/rewards/sarm/sarm_utils.py diff --git a/src/lerobot/rewards/sarm/__init__.py b/src/lerobot/rewards/sarm/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/lerobot/rewards/sarm/compute_rabc_weights.py b/src/lerobot/rewards/sarm/compute_rabc_weights.py new file mode 100644 index 000000000..3a79d3fff --- /dev/null +++ b/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 /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 /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 /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 /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 /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() diff --git a/src/lerobot/rewards/sarm/configuration_sarm.py b/src/lerobot/rewards/sarm/configuration_sarm.py new file mode 100644 index 000000000..2a05b9b2b --- /dev/null +++ b/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 diff --git a/src/lerobot/rewards/sarm/modeling_sarm.py b/src/lerobot/rewards/sarm/modeling_sarm.py new file mode 100644 index 000000000..f920607bc --- /dev/null +++ b/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) diff --git a/src/lerobot/rewards/sarm/processor_sarm.py b/src/lerobot/rewards/sarm/processor_sarm.py new file mode 100644 index 000000000..d60914c4a --- /dev/null +++ b/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, + ), + ) diff --git a/src/lerobot/rewards/sarm/sarm_utils.py b/src/lerobot/rewards/sarm/sarm_utils.py new file mode 100644 index 000000000..e7231db2e --- /dev/null +++ b/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])