robometer/rewardeval_ui
4
1#!/usr/bin/env python32"""3PrefSampler class for producing batches of preference data.4"""5 6from typing import Dict, List, Optional, Any7 8import random9 10from robometer.data.dataset_types import PreferenceSample, Trajectory11from robometer.data.samplers.base import RBMBaseSampler12from robometer.data.datasets.helpers import (13 DataGenStrat,14 convert_continuous_to_discrete_bins,15)16from robometer.utils.logger import get_logger, rank_0_info, trace17from robometer.utils.timer import timer18 19logger = get_logger()20 21 22class PrefSampler(RBMBaseSampler):23 """Data generator for producing batches of preference prediction data."""24 25 def __init__(self, is_evaluation=False, **kwargs):26 super().__init__(**kwargs)27 28 self.dataset_preference_ratio = self.config.dataset_preference_ratio29 self.preference_strategy_ratio: List[float] = self.config.preference_strategy_ratio30 self._has_suboptimal = (31 any(len(indices) > 0 for indices in self.suboptimal_by_task.values()) if self.suboptimal_by_task else False32 )33 rank_0_info(f"[PREF SAMPLER] Has suboptimal: {self._has_suboptimal}")34 35 # Initialize preference dataset36 self._load_preference_dataset()37 38 def _generate_sample(self, item: dict, preferred_strategy: Optional[DataGenStrat] = None):39 """Generate a preference sample from an item.40 41 If the item has a non-successful quality label, it will be used as the rejected42 trajectory and an optimal trajectory from the same task will be found as the chosen one.43 Otherwise, normal preference sampling logic is used.44 45 Args:46 item: The trajectory item47 preferred_strategy: Optional strategy to use (if None, will select strategy based on ratios)48 """49 quality_label = item["quality_label"]50 use_partial_success = item.get("partial_success") is not None51 52 # Handle non-successful trajectories: use as rejected, find optimal from same task as chosen53 # skip this for trajectories with partial_success which we will handle with partial success logic54 if quality_label != "successful" and not use_partial_success:55 traj_id = item["id"]56 task_name = item["task"]57 58 logger.trace(59 f"[PREF SAMPLER] Non-successful quality detected for ID={traj_id}, using as rejected trajectory, task={task_name}"60 )61 62 # Find optimal trajectories from the same task63 same_task_optimal_indices = self.optimal_by_task.get(task_name, [])64 65 if not same_task_optimal_indices:66 logger.trace(67 f"[PREF SAMPLER] No optimal trajectories found for task '{task_name}', falling through to normal sampling"68 )69 return self._create_pref_sample(item, preferred_strategy=preferred_strategy)70 71 # Select a random optimal trajectory from the same task as chosen72 chosen_idx = random.choice(same_task_optimal_indices)73 chosen_traj_dict = self.dataset[chosen_idx]74 75 chosen_trajectory = self._get_traj_from_data(chosen_traj_dict)76 rejected_trajectory = self._get_traj_from_data(item)77 78 sample = PreferenceSample(79 chosen_trajectory=chosen_trajectory,80 rejected_trajectory=rejected_trajectory,81 data_gen_strategy=DataGenStrat.SUBOPTIMAL.value,82 )83 84 logger.trace(85 f"[PREF SAMPLER] Created preference sample for non-successful traj ID={traj_id} with optimal traj from same task"86 )87 return sample88 89 return self._create_pref_sample(item, preferred_strategy=preferred_strategy)90 91 def _execute_strategy(92 self, strategy: DataGenStrat, chosen_traj: Dict[str, Any], use_partial_success: bool93 ) -> tuple[Dict[str, Any], str, Dict[str, Any]] | None:94 """Execute a strategy to get rejected trajectory.95 96 Args:97 strategy: The strategy to execute98 chosen_traj: The chosen trajectory99 use_partial_success: Whether this trajectory uses partial_success100 101 Returns:102 Tuple of (rejected_traj, rejected_subsample_strategy, chosen_traj) or None if failed103 Note: chosen_traj may be swapped with rejected_traj for partial_success trajectories104 """105 max_retries = 3106 rejected_subsample_strategy = None107 rejected_traj = None108 109 if strategy == DataGenStrat.REWIND:110 rejected_traj = chosen_traj.copy()111 rejected_subsample_strategy = "subsample_rewind"112 elif strategy == DataGenStrat.SUBOPTIMAL:113 for _ in range(max_retries):114 rejected_traj = self._get_same_task_suboptimal(chosen_traj)115 if rejected_traj is not None:116 # For trajectories with partial_success, if the returned trajectory has higher partial_success, swap them117 if use_partial_success:118 chosen_partial_success = chosen_traj.get("partial_success")119 rejected_partial_success = rejected_traj.get("partial_success")120 if rejected_partial_success is not None and chosen_partial_success is not None:121 if rejected_partial_success > chosen_partial_success:122 logger.trace(123 f"[PREF SAMPLER] Swapping trajectories: found higher partial_success "124 f"({rejected_partial_success} > {chosen_partial_success})"125 )126 rejected_traj, chosen_traj = chosen_traj, rejected_traj127 break128 rejected_subsample_strategy = "subsample_forward"129 elif strategy == DataGenStrat.DIFFERENT_TASK:130 for _ in range(max_retries):131 rejected_traj = self._get_different_video_traj(chosen_traj)132 if rejected_traj is not None:133 break134 rejected_subsample_strategy = "subsample_forward"135 elif strategy == DataGenStrat.REVERSE_PROGRESS:136 rejected_traj = chosen_traj.copy()137 rejected_subsample_strategy = "subsample_reverse"138 else:139 return None140 141 if rejected_traj is None:142 return None143 144 return (rejected_traj, rejected_subsample_strategy, chosen_traj)145 146 def _create_pref_sample_from_dataset(self) -> PreferenceSample:147 """Create a preference sample from the loaded preference dataset."""148 if not self.preferences:149 return None150 151 # For now, return a simple preference sample152 # This can be enhanced later when we have actual preference data153 random.choice(self.preferences)154 155 # This is a placeholder - would need to be implemented based on actual preference data structure156 return None157 158 def _load_preference_dataset(self):159 """Load the preference dataset from disk or hub if provided."""160 self.preferences = []161 162 # For now, we'll use empty preferences since the config structure has changed163 # This can be updated later if needed164 rank_0_info("[PREF SAMPLER] No preference dataset provided, will use random sampling for preferences")165 return166 167 def _create_preference_sample(self) -> PreferenceSample:168 """Create a preference prediction sample: chosen vs rejected where chosen is preferred.169 Either from dataset or from generated trajectories.170 171 Returns:172 PreferenceSample: A preference sample with chosen (preferred) vs rejected173 (suboptimal) trajectories and associated metadata174 """175 176 with timer("create_preference_sample", verbose=False):177 if random.random() < self.dataset_preference_ratio and self.preferences:178 # Use preference trajectories from dataset179 return self._create_pref_sample_from_dataset()180 else:181 return self._create_pref_sample()182 183 def _create_pref_sample(184 self, chosen_traj: Optional[Dict[str, Any]] = None, preferred_strategy: Optional[DataGenStrat] = None185 ) -> PreferenceSample:186 """Create a preference prediction sample using various rejected trajectory generation strategies.187 188 Rewind Same Task189 - Creates a suboptimal trajectory by rewinding the chosen trajectory190 191 Suboptimal/Failure Same Task192 - Uses existing suboptimal/failure trajectories from the same task193 194 Different Task195 - Uses trajectories from completely different tasks196 197 Returns:198 PreferenceSample: A preference sample with chosen (preferred) vs rejected199 (suboptimal) trajectories and associated metadata200 201 Raises:202 ValueError: If no chosen trajectories are available for preference generation203 RuntimeError: If all strategies fail and fallback rewind also fails204 """205 # Log when preference sampler is called206 traj_id = chosen_traj["id"] if chosen_traj is not None else "sampling_new"207 logger.trace(f"[PREF SAMPLER] Creating preference sample for trajectory ID: {traj_id}")208 209 # Use provided chosen trajectory if given; otherwise sample one210 if chosen_traj is None:211 # Use preprocessed chosen trajectories from index maps212 if not self.optimal_by_task:213 return None214 215 # Filter out tasks with empty optimal_indices to avoid infinite loop216 valid_tasks = {217 task: indices218 for task, indices in self.optimal_by_task.items()219 if indices # Only include tasks with non-empty indices220 }221 222 if not valid_tasks:223 # No valid tasks with optimal trajectories available224 return None225 226 # Get a random task and chosen trajectory from it227 task_name = random.choice(list(valid_tasks.keys()))228 optimal_indices = valid_tasks[task_name]229 230 # Double-check that we have valid indices (should always be true now)231 if not optimal_indices:232 return None233 234 chosen_idx = random.choice(optimal_indices)235 chosen_traj = self.dataset[chosen_idx]236 237 # Initialize variables for strategy selection238 rejected_traj = None239 strategy_used = None240 rejected_subsample_strategy = None241 242 # Check if this trajectory uses partial_success243 use_partial_success = chosen_traj.get("partial_success") is not None244 if use_partial_success:245 partial_success = chosen_traj.get("partial_success")246 logger.trace(247 f"[PREF SAMPLER] Trajectory with partial_success detected (ID: {chosen_traj.get('id', 'unknown')}, partial_success: {partial_success})"248 )249 250 # Strategy selection: use preferred_strategy if provided, otherwise select based on ratios251 if preferred_strategy is not None:252 # Use the preferred strategy directly253 logger.trace(f"[PREF SAMPLER] Using preferred strategy: {preferred_strategy.value}")254 result = self._execute_strategy(preferred_strategy, chosen_traj, use_partial_success)255 if result is None:256 logger.trace(f"[PREF SAMPLER] Preferred strategy {preferred_strategy.value} failed, returning None")257 return None258 rejected_traj, rejected_subsample_strategy, chosen_traj = result259 strategy_used = preferred_strategy260 attempt = 1 # Set attempt for preferred strategy path261 else:262 # Strategy selection with rebalancing on failure263 strategies = []264 if self.preference_strategy_ratio[0] > 0:265 strategies.append((DataGenStrat.REWIND, self.preference_strategy_ratio[0]))266 if self._has_suboptimal and self.preference_strategy_ratio[1] > 0:267 strategies.append((DataGenStrat.SUBOPTIMAL, self.preference_strategy_ratio[1]))268 if self.preference_strategy_ratio[2] > 0:269 strategies.append((DataGenStrat.DIFFERENT_TASK, self.preference_strategy_ratio[2]))270 if self.preference_strategy_ratio[3] > 0:271 strategies.append((DataGenStrat.REVERSE_PROGRESS, self.preference_strategy_ratio[3]))272 273 max_attempts = 10 # Limit retry attempts to prevent infinite loops274 max_strategy_attempts = 3 # Maximum attempts per strategy before removing it275 attempt = 0276 277 # Track attempts per strategy278 strategy_attempt_counts = {strat: 0 for strat, _ in strategies}279 280 while rejected_traj is None and attempt < max_attempts:281 attempt += 1282 283 # Check if we have any strategies left284 if not strategies:285 return None286 287 # Rebalance probabilities based on remaining strategies288 total_prob = sum(prob for _, prob in strategies)289 if total_prob == 0:290 return None291 292 # Normalize probabilities293 normalized_strategies = [(strat, prob / total_prob) for strat, prob in strategies]294 295 # Select strategy based on rebalanced probabilities296 prob = random.random()297 cumulative_prob = 0.0298 selected_strategy = None299 300 for strat, normalized_prob in normalized_strategies:301 cumulative_prob += normalized_prob302 if prob <= cumulative_prob:303 selected_strategy = strat304 break305 306 # Log strategy attempt307 logger.trace(308 f"[PREF SAMPLER] Attempt {attempt}/{max_attempts}: Trying strategy {selected_strategy.value if selected_strategy else 'None'}"309 )310 311 # Execute selected strategy312 result = self._execute_strategy(selected_strategy, chosen_traj, use_partial_success)313 if result is not None:314 rejected_traj, rejected_subsample_strategy, chosen_traj = result315 strategy_used = selected_strategy316 logger.trace(f"[PREF SAMPLER] Strategy {selected_strategy.value} succeeded on attempt {attempt}")317 else:318 # Strategy failed - increment attempt count319 strategy_attempt_counts[selected_strategy] = strategy_attempt_counts.get(selected_strategy, 0) + 1320 failed_count = strategy_attempt_counts[selected_strategy]321 322 logger.trace(323 f"[PREF SAMPLER] Strategy {selected_strategy.value} failed (failure count: {failed_count}/{max_strategy_attempts})"324 )325 326 # Only remove strategy if it has failed max_strategy_attempts times327 if strategy_attempt_counts[selected_strategy] >= max_strategy_attempts:328 logger.trace(329 f"[PREF SAMPLER] Removing strategy {selected_strategy.value} after {max_strategy_attempts} consecutive failures"330 )331 strategies = [(strat, prob) for strat, prob in strategies if strat != selected_strategy]332 continue333 334 # If we still don't have a sample after all attempts, return None335 if rejected_traj is None:336 logger.trace(337 f"[PREF SAMPLER] Failed to generate preference sample after {max_attempts} attempts - all strategies exhausted"338 )339 return None340 341 chosen_subsample_strategy = "subsample_forward"342 chosen_trajectory = self._get_traj_from_data(chosen_traj, subsample_strategy=chosen_subsample_strategy)343 344 rejected_trajectory = self._get_traj_from_data(rejected_traj, subsample_strategy=rejected_subsample_strategy)345 346 if rejected_trajectory is None or chosen_trajectory is None:347 return None348 349 # If our strategy is different task, make sure the rejected trajectory has 0 progress and 0 success labels350 if strategy_used in [351 DataGenStrat.DIFFERENT_TASK,352 DataGenStrat.DIFFERENT_TASK_INSTRUCTION,353 ]:354 rejected_trajectory.target_progress = [0.0] * len(rejected_trajectory.target_progress)355 if self.config.progress_loss_type.lower() == "discrete":356 rejected_trajectory.target_progress = convert_continuous_to_discrete_bins(357 rejected_trajectory.target_progress, self.config.progress_discrete_bins358 )359 # Also set success labels to 0.0 (predict 0 success for different task trajectories)360 if rejected_trajectory.success_label is not None:361 rejected_trajectory.success_label = [0.0] * len(rejected_trajectory.success_label)362 363 # Create preference sample structure364 sample = PreferenceSample(365 chosen_trajectory=chosen_trajectory,366 rejected_trajectory=rejected_trajectory,367 data_gen_strategy=strategy_used.value,368 )369 sample.resample_attempts = attempt370 return sample371