From 9ebc144b300e329b32deb6bf947c14a714193a2d Mon Sep 17 00:00:00 2001 From: Michel Aractingi Date: Mon, 12 Jan 2026 11:39:01 +0100 Subject: [PATCH] linter + missing files --- src/lerobot/configs/train.py | 17 +- src/lerobot/policies/sarm/rabc.py | 1 - src/lerobot/scripts/lerobot_train.py | 87 ++++---- src/lerobot/utils/rabc.py | 288 -------------------------- src/lerobot/utils/sample_weighting.py | 6 +- 5 files changed, 39 insertions(+), 360 deletions(-) delete mode 100644 src/lerobot/utils/rabc.py diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index 7a5eee77d..c8c1d6791 100644 --- a/src/lerobot/configs/train.py +++ b/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`""" diff --git a/src/lerobot/policies/sarm/rabc.py b/src/lerobot/policies/sarm/rabc.py index 42b11c70d..3fdbe4eda 100644 --- a/src/lerobot/policies/sarm/rabc.py +++ b/src/lerobot/policies/sarm/rabc.py @@ -306,4 +306,3 @@ class RABCWeights: "delta_std": self.delta_std, "kappa": self.kappa, } - diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 286c69906..bc8cc79e8 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/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() diff --git a/src/lerobot/utils/rabc.py b/src/lerobot/utils/rabc.py deleted file mode 100644 index dc0c61c69..000000000 --- a/src/lerobot/utils/rabc.py +++ /dev/null @@ -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, - } diff --git a/src/lerobot/utils/sample_weighting.py b/src/lerobot/utils/sample_weighting.py index ec3b5ab07..26303ee5e 100644 --- a/src/lerobot/utils/sample_weighting.py +++ b/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"} -