linter + missing files

src/lerobot/configs/train.py

@@ -29,6 +29,7 @@ from lerobot.configs.policies import PreTrainedConfig
 from lerobot.optim import OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
 from lerobot.utils.hub import HubMixin
+from lerobot.utils.sample_weighting import SampleWeightingConfig
 
 TRAIN_CONFIG_NAME = "train_config.json"
 
@@ -67,12 +68,8 @@ class TrainPipelineConfig(HubMixin):
     wandb: WandBConfig = field(default_factory=WandBConfig)
     peft: PeftConfig | None = None
 
-    # RA-BC (Reward-Aligned Behavior Cloning) parameters
-    use_rabc: bool = False  # Enable reward-weighted training
-    rabc_progress_path: str | None = None  # Path to precomputed SARM progress parquet file
-    rabc_kappa: float = 0.01  # Hard threshold for high-quality samples
-    rabc_epsilon: float = 1e-6  # Small constant for numerical stability
-    rabc_head_mode: str | None = "sparse"  # For dual-head models: "sparse" or "dense"
+    # Sample weighting configuration (e.g., for RA-BC training)
+    sample_weighting: SampleWeightingConfig | None = None
 
     # Rename map for the observation to override the image and state keys
     rename_map: dict[str, str] = field(default_factory=dict)
@@ -140,14 +137,6 @@ class TrainPipelineConfig(HubMixin):
                 "'policy.repo_id' argument missing. Please specify it to push the model to the hub."
             )
 
-        if self.use_rabc and not self.rabc_progress_path:
-            # Auto-detect from dataset path
-            repo_id = self.dataset.repo_id
-            if self.dataset.root:
-                self.rabc_progress_path = str(Path(self.dataset.root) / "sarm_progress.parquet")
-            else:
-                self.rabc_progress_path = f"hf://datasets/{repo_id}/sarm_progress.parquet"
-
     @classmethod
     def __get_path_fields__(cls) -> list[str]:
         """This enables the parser to load config from the policy using `--policy.path=local/dir`"""

src/lerobot/policies/sarm/rabc.py

@@ -306,4 +306,3 @@ class RABCWeights:
             "delta_std": self.delta_std,
             "kappa": self.kappa,
         }
-

src/lerobot/scripts/lerobot_train.py

@@ -63,7 +63,7 @@ def update_policy(
     accelerator: Accelerator,
     lr_scheduler=None,
     lock=None,
-    rabc_weights_provider=None,
+    sample_weighter=None,
 ) -> tuple[MetricsTracker, dict]:
     """
     Performs a single training step to update the policy's weights.
@@ -80,7 +80,7 @@ def update_policy(
         accelerator: The Accelerator instance for distributed training and mixed precision.
         lr_scheduler: An optional learning rate scheduler.
         lock: An optional lock for thread-safe optimizer updates.
-        rabc_weights_provider: Optional RABCWeights instance for sample weighting.
+        sample_weighter: Optional SampleWeighter instance for per-sample loss weighting.
 
     Returns:
         A tuple containing:
@@ -90,27 +90,30 @@ def update_policy(
     start_time = time.perf_counter()
     policy.train()
 
-    # Get RA-BC weights if enabled
-    rabc_batch_weights = None
-    rabc_batch_stats = None
-    if rabc_weights_provider is not None:
-        rabc_batch_weights, rabc_batch_stats = rabc_weights_provider.compute_batch_weights(batch)
+    # Compute sample weights if a weighter is provided
+    sample_weights = None
+    weight_stats = None
+    if sample_weighter is not None:
+        sample_weights, weight_stats = sample_weighter.compute_batch_weights(batch)
 
     # Let accelerator handle mixed precision
     with accelerator.autocast():
-        # Use per-sample loss when RA-BC is enabled for proper weighting
-        if rabc_batch_weights is not None:
-            # Get per-sample losses
-            per_sample_loss, output_dict = policy.forward(batch, reduction="none")
+        if sample_weights is not None:
+            # Use per-sample loss for weighted training
+            # Note: Policies supporting sample weighting must implement forward(batch, reduction="none")
+            per_sample_loss, output_dict = policy.forward(batch, reduction="none")  # type: ignore[call-arg]
 
-            # Apply RA-BC weights: L_RA-BC = Σ(w_i * l_i) / (Σw_i + ε)
-            # rabc_batch_weights is already normalized to sum to batch_size
+            # Apply sample weights: L_weighted = Σ(w_i * l_i) / (Σw_i + ε)
+            # Weights are already normalized to sum to batch_size
             epsilon = 1e-6
-            loss = (per_sample_loss * rabc_batch_weights).sum() / (rabc_batch_weights.sum() + epsilon)
-            # Log raw mean weight (before normalization) - this is the meaningful metric
-            output_dict["rabc_mean_weight"] = rabc_batch_stats["raw_mean_weight"]
-            output_dict["rabc_num_zero_weight"] = rabc_batch_stats["num_zero_weight"]
-            output_dict["rabc_num_full_weight"] = rabc_batch_stats["num_full_weight"]
+            loss = (per_sample_loss * sample_weights).sum() / (sample_weights.sum() + epsilon)
+
+            # Log weighting statistics (weight_stats is set when sample_weights is not None)
+            if output_dict is None:
+                output_dict = {}
+            if weight_stats is not None:
+                for key, value in weight_stats.items():
+                    output_dict[f"sample_weight_{key}"] = value
         else:
             loss, output_dict = policy.forward(batch)
 
@@ -142,10 +145,10 @@ def update_policy(
         accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update()
 
     train_metrics.loss = loss.item()
-    train_metrics.grad_norm = grad_norm.item()
+    train_metrics.grad_norm = grad_norm.item() if hasattr(grad_norm, "item") else float(grad_norm)
     train_metrics.lr = optimizer.param_groups[0]["lr"]
     train_metrics.update_s = time.perf_counter() - start_time
-    return train_metrics, output_dict
+    return train_metrics, output_dict if output_dict is not None else {}
 
 
 def get_default_peft_configuration(policy_type):
@@ -366,28 +369,14 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
         logging.info("Creating optimizer and scheduler")
     optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
 
-    # Load precomputed SARM progress for RA-BC if enabled
-    # Generate progress using: src/lerobot/policies/sarm/compute_rabc_weights.py
-    rabc_weights = None
-    if cfg.use_rabc:
-        from lerobot.utils.rabc import RABCWeights
+    # Create sample weighter if configured (e.g., for RA-BC training)
+    sample_weighter = None
+    if cfg.sample_weighting is not None:
+        from lerobot.utils.sample_weighting import make_sample_weighter
 
-        # Get chunk_size from policy config
-        chunk_size = getattr(policy.config, "chunk_size", None)
-        if chunk_size is None:
-            raise ValueError("Chunk size is not found in policy config")
-
-        head_mode = getattr(cfg, "rabc_head_mode", "sparse")
-        logging.info(f"Loading SARM progress for RA-BC from {cfg.rabc_progress_path}")
-        logging.info(f"Using chunk_size={chunk_size} from policy config, head_mode={head_mode}")
-        rabc_weights = RABCWeights(
-            progress_path=cfg.rabc_progress_path,
-            chunk_size=chunk_size,
-            head_mode=head_mode,
-            kappa=getattr(cfg, "rabc_kappa", 0.01),
-            epsilon=getattr(cfg, "rabc_epsilon", 1e-6),
-            device=device,
-        )
+        if is_main_process:
+            logging.info(f"Creating sample weighter: {cfg.sample_weighting.type}")
+        sample_weighter = make_sample_weighter(cfg.sample_weighting, policy, device)
 
     step = 0  # number of policy updates (forward + backward + optim)
 
@@ -486,7 +475,7 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
             cfg.optimizer.grad_clip_norm,
             accelerator=accelerator,
             lr_scheduler=lr_scheduler,
-            rabc_weights_provider=rabc_weights,
+            sample_weighter=sample_weighter,
         )
 
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
@@ -503,16 +492,10 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
                 wandb_log_dict = train_tracker.to_dict()
                 if output_dict:
                     wandb_log_dict.update(output_dict)
-                # Log RA-BC statistics if enabled
-                if rabc_weights is not None:
-                    rabc_stats = rabc_weights.get_stats()
-                    wandb_log_dict.update(
-                        {
-                            "rabc_delta_mean": rabc_stats["delta_mean"],
-                            "rabc_delta_std": rabc_stats["delta_std"],
-                            "rabc_num_frames": rabc_stats["num_frames"],
-                        }
-                    )
+                # Log sample weighting statistics if enabled
+                if sample_weighter is not None:
+                    weighter_stats = sample_weighter.get_stats()
+                    wandb_log_dict.update({f"sample_weighting/{k}": v for k, v in weighter_stats.items()})
                 wandb_logger.log_dict(wandb_log_dict, step)
             train_tracker.reset_averages()
 

src/lerobot/utils/rabc.py

@@ -1,288 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import logging
-from pathlib import Path
-
-import numpy as np
-import pandas as pd
-import torch
-from huggingface_hub import hf_hub_download
-
-
-def resolve_hf_path(path: str | Path) -> Path:
-    """Resolve a path that may be a HuggingFace URL (hf://datasets/...) to a local path."""
-    path_str = str(path)
-    if path_str.startswith("hf://datasets/"):
-        parts = path_str.replace("hf://datasets/", "").split("/")
-        repo_id = "/".join(parts[:2])
-        filename = "/".join(parts[2:])
-        return Path(hf_hub_download(repo_id=repo_id, filename=filename, repo_type="dataset"))
-    return Path(path)
-
-
-class RABCWeights:
-    """
-    Load precomputed SARM progress values and compute RA-BC weights during training.
-
-    Progress values are loaded from a parquet file (generated by compute_rabc_weights.py).
-    During training, computes:
-        - progress_delta = progress[t + chunk_size] - progress[t]
-        - rabc_weight based on the delta (paper Eq. 8-9)
-
-    Args:
-        progress_path: Path to parquet file with precomputed progress values
-        chunk_size: Number of frames ahead for computing progress delta
-        head_mode: Which SARM head to use ("sparse" or "dense")
-        kappa: Hard threshold for high-quality samples (default: 0.01)
-        epsilon: Small constant for numerical stability (default: 1e-6)
-        fallback_weight: Weight to use for frames without valid delta (default: 1.0)
-        device: Device to return tensors on
-    """
-
-    def __init__(
-        self,
-        progress_path: str | Path,
-        chunk_size: int = 50,
-        head_mode: str = "sparse",
-        kappa: float = 0.01,
-        epsilon: float = 1e-6,
-        fallback_weight: float = 1.0,
-        device: torch.device = None,
-    ):
-        self.progress_path = resolve_hf_path(progress_path)
-        self.chunk_size = chunk_size
-        self.head_mode = head_mode
-        self.kappa = kappa
-        self.epsilon = epsilon
-        self.fallback_weight = fallback_weight
-        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-        # Determine progress column name
-        self.progress_column = f"progress_{head_mode}"
-
-        # Load progress values
-        logging.info(f"Loading SARM progress values from {self.progress_path}")
-        self.df = pd.read_parquet(self.progress_path)
-
-        # Check if the requested head mode column exists
-        if self.progress_column not in self.df.columns:
-            available = [c for c in self.df.columns if c.startswith("progress")]
-            raise ValueError(
-                f"Column '{self.progress_column}' not found. Available progress columns: {available}"
-            )
-
-        logging.info(f"Using progress column: {self.progress_column}")
-
-        self.progress_lookup = {}
-        self.episode_lookup = {}
-
-        for _, row in self.df.iterrows():
-            global_idx = int(row["index"])
-            progress = row[self.progress_column]
-            episode_idx = int(row["episode_index"])
-
-            if not np.isnan(progress):
-                self.progress_lookup[global_idx] = float(progress)
-            self.episode_lookup[global_idx] = episode_idx
-
-        # Build episode boundaries for delta computation
-        self.episode_boundaries = {}
-        for episode_idx in self.df["episode_index"].unique():
-            ep_df = self.df[self.df["episode_index"] == episode_idx]
-            self.episode_boundaries[int(episode_idx)] = {
-                "start": int(ep_df["index"].min()),
-                "end": int(ep_df["index"].max()) + 1,
-            }
-
-        logging.info(f"Loaded {len(self.progress_lookup)} frame progress values")
-        logging.info(f"Chunk size for delta computation: {chunk_size}")
-
-        # Compute global statistics for weight computation
-        self._compute_global_stats()
-
-    def _compute_global_stats(self):
-        """Compute global mean and std of progress deltas for weight calculation."""
-        all_deltas = []
-
-        for global_idx, progress in self.progress_lookup.items():
-            episode_idx = self.episode_lookup.get(global_idx)
-            if episode_idx is None:
-                continue
-
-            bounds = self.episode_boundaries.get(episode_idx)
-            if bounds is None:
-                continue
-
-            future_idx = global_idx + self.chunk_size
-            if future_idx >= bounds["end"]:
-                # Near end of episode: use last frame's progress
-                future_idx = bounds["end"] - 1
-
-            future_progress = self.progress_lookup.get(future_idx)
-            if future_progress is not None:
-                delta = future_progress - progress
-                all_deltas.append(delta)
-
-        if all_deltas:
-            self.delta_mean = max(np.mean(all_deltas), 0.0)
-            self.delta_std = max(np.std(all_deltas), self.epsilon)
-            logging.info(f"Progress delta stats: mean={self.delta_mean:.4f}, std={self.delta_std:.4f}")
-        else:
-            self.delta_mean = 0.0
-            self.delta_std = self.epsilon
-            logging.warning("No valid progress deltas found, using default stats")
-
-    def compute_batch_weights(self, batch: dict) -> tuple[torch.Tensor, dict]:
-        """
-        Compute RA-BC weights for a batch.
-
-        For each sample:
-        1. Get progress at current frame
-        2. Get progress at frame + chunk_size (within same episode)
-        3. Compute delta = future_progress - current_progress
-        4. Compute weight using paper Eq. 8-9
-
-        Args:
-            batch: Training batch containing "index" key with global frame indices
-
-        Returns:
-            Tuple of:
-            - Weights tensor (batch_size,) normalized to sum to batch_size
-            - Stats dict with raw_mean_weight, num_zero_weight, num_full_weight
-        """
-        indices = batch.get("index")
-        if indices is None:
-            logging.warning("RA-BC: Batch missing 'index' key, using uniform weights")
-            batch_size = self._get_batch_size(batch)
-            return torch.ones(batch_size, device=self.device), {"raw_mean_weight": 1.0}
-
-        # Convert to list of ints
-        if isinstance(indices, torch.Tensor):
-            indices = indices.cpu().numpy().tolist()
-        elif isinstance(indices, np.ndarray):
-            indices = indices.tolist()
-
-        # Compute deltas and weights for each sample
-        deltas = []
-        for idx in indices:
-            idx = int(idx)
-            delta = self._compute_delta(idx)
-            deltas.append(delta)
-
-        deltas = np.array(deltas, dtype=np.float32)
-
-        # Compute weights from deltas
-        weights = self._compute_weights(deltas)
-
-        # Compute stats before normalization for logging
-        raw_mean_weight = float(np.nanmean(weights))
-        num_zero_weight = int(np.sum(weights == 0))
-        num_full_weight = int(np.sum(weights == 1.0))
-        batch_stats = {
-            "raw_mean_weight": raw_mean_weight,
-            "num_zero_weight": num_zero_weight,
-            "num_full_weight": num_full_weight,
-        }
-
-        weights = torch.tensor(weights, device=self.device, dtype=torch.float32)
-
-        # Normalize to sum to batch_size
-        batch_size = len(weights)
-        weight_sum = weights.sum() + self.epsilon
-        weights = weights * batch_size / weight_sum
-
-        return weights, batch_stats
-
-    def _compute_delta(self, global_idx: int) -> float:
-        """Compute progress delta for a single frame."""
-        current_progress = self.progress_lookup.get(global_idx)
-        if current_progress is None:
-            return np.nan
-
-        episode_idx = self.episode_lookup.get(global_idx)
-        if episode_idx is None:
-            return np.nan
-
-        bounds = self.episode_boundaries.get(episode_idx)
-        if bounds is None:
-            return np.nan
-
-        future_idx = global_idx + self.chunk_size  # Δ = chunk_size
-        if future_idx >= bounds["end"]:
-            # Near end of episode: use last frame's progress instead
-            future_idx = bounds["end"] - 1
-
-        future_progress = self.progress_lookup.get(future_idx)
-        if future_progress is None:
-            return np.nan
-
-        return future_progress - current_progress
-
-    def _compute_weights(self, deltas: np.ndarray) -> np.ndarray:
-        """
-        Compute RA-BC weights from progress deltas.
-
-        Following paper Eq. 8-9:
-        - Soft weight: ˜wi = clip((ri − (µ − 2σ)) / (4σ + ε), 0, 1)
-        - Final weight: wi = 1{ri > κ} + 1{0 ≤ ri ≤ κ}˜wi
-
-        Returns:
-            Array of weights
-        """
-        valid_mask = ~np.isnan(deltas)
-
-        # Compute soft weights using global statistics
-        lower_bound = self.delta_mean - 2 * self.delta_std
-        soft_weights = (deltas - lower_bound) / (4 * self.delta_std + self.epsilon)
-        soft_weights = np.clip(soft_weights, 0.0, 1.0)
-
-        # Apply paper's Eq. 9
-        weights = np.zeros_like(deltas, dtype=np.float32)
-
-        # High quality: ri > kappa → weight = 1
-        high_quality_mask = deltas > self.kappa
-        weights[high_quality_mask] = 1.0
-
-        # Moderate quality: 0 <= ri <= kappa → weight = soft_weight
-        moderate_mask = (deltas >= 0) & (deltas <= self.kappa)
-        weights[moderate_mask] = soft_weights[moderate_mask]
-
-        # Negative progress: ri < 0 → weight = 0 (already 0)
-        # Invalid (NaN): use fallback weight
-        weights[~valid_mask] = self.fallback_weight
-
-        return weights
-
-    def _get_batch_size(self, batch: dict) -> int:
-        """Determine batch size from batch."""
-        for key in ["action", "index"]:
-            if key in batch:
-                val = batch[key]
-                if isinstance(val, (torch.Tensor, np.ndarray)):
-                    return val.shape[0]
-        return 1
-
-    def get_stats(self) -> dict:
-        """Get statistics."""
-        return {
-            "num_frames": len(self.progress_lookup),
-            "chunk_size": self.chunk_size,
-            "head_mode": self.head_mode,
-            "delta_mean": self.delta_mean,
-            "delta_std": self.delta_std,
-            "kappa": self.kappa,
-        }

src/lerobot/utils/sample_weighting.py

@@ -140,10 +140,7 @@ def make_sample_weighter(
         # No-op weighter that returns uniform weights
         return UniformWeighter(device=device)
 
-    raise ValueError(
-        f"Unknown sample weighting type: '{config.type}'. "
-        f"Supported types: 'rabc', 'uniform'"
-    )
+    raise ValueError(f"Unknown sample weighting type: '{config.type}'. Supported types: 'rabc', 'uniform'")
 
 
 def _make_rabc_weighter(
@@ -210,4 +207,3 @@ class UniformWeighter:
     def get_stats(self) -> dict:
         """Return empty stats for uniform weighting."""
         return {"type": "uniform"}
-