robometer/rewardeval_ui

1from typing import Dict, Any, Optional2 3import random4import torch5 6from robometer.data.dataset_types import ProgressSample, Trajectory7from robometer.data.samplers.base import RBMBaseSampler8from robometer.data.datasets.helpers import (9    DataGenStrat,10    load_embeddings_from_path,11    convert_continuous_to_discrete_bins,12)13from robometer.utils.distributed import rank_0_print14from robometer.utils.logger import get_logger15 16logger = get_logger()17 18 19class ProgressSampler(RBMBaseSampler):20    """Data generator for progress samples."""21 22    def __init__(self, is_evaluation=False, **kwargs):23        super().__init__(**kwargs)24 25    def _generate_sample(self, item: Dict[str, Any], preferred_strategy: Optional[DataGenStrat] = None):26        return self._create_progress_sample(item, preferred_strategy=preferred_strategy)27 28    def _execute_strategy(self, strategy: DataGenStrat, traj: Dict[str, Any]) -> tuple[Dict[str, Any], str] | None:29        """Execute a strategy to get processed trajectory.30 31        Args:32            strategy: The strategy to execute33            traj: The trajectory to process34 35        Returns:36            Tuple of (processed_traj, subsample_strategy) or None if failed37        """38        if strategy == DataGenStrat.FORWARD_PROGRESS:39            return (traj, "subsample_forward")40        elif strategy == DataGenStrat.REVERSE_PROGRESS:41            return (traj, "subsample_reverse")42        elif strategy == DataGenStrat.REWIND:43            return (traj, "subsample_rewind")44        elif strategy == DataGenStrat.DIFFERENT_TASK_INSTRUCTION:45            processed_traj = self._get_different_task_instruction(traj)46            if processed_traj is None:47                return None48            return (processed_traj, "subsample_forward")49        else:50            return None51 52    def _create_progress_sample(self, traj: Dict[str, Any], preferred_strategy: Optional[DataGenStrat] = None):53        """Create a progress sample using normalized and rebalanced strategy selection.54 55        Implements four strategies:56        1. Different Task: Use trajectory from different task (progress set to 0.0)57        2. Forward Progress: Sample with forward direction (start < middle < end)58        3. Reverse Progress: Sample with reverse direction (end < middle < start)59        4. Rewind: Sample with rewind direction (start < end < middle)60        """61        # Initialize variables for strategy selection62        processed_traj = None63        strategy_used = None64        subsample_strategy = None65 66        # Strategy selection: use preferred_strategy if provided, otherwise select based on ratios67        if preferred_strategy is not None:68            # Use the preferred strategy directly69            logger.trace(f"[PROGRESS SAMPLER] Using preferred strategy: {preferred_strategy.value}")70            result = self._execute_strategy(preferred_strategy, traj)71            if result is None:72                logger.trace(f"[PROGRESS SAMPLER] Preferred strategy {preferred_strategy.value} failed, returning None")73                return None74            processed_traj, subsample_strategy = result75            strategy_used = preferred_strategy76            attempt = 1  # Set attempt for preferred strategy path77        else:78            # Strategy setup with rebalancing on failure79            # [different_task_instruction, forward_progress, reverse_progress, rewind]80            strategies = [81                (82                    DataGenStrat.DIFFERENT_TASK_INSTRUCTION,83                    self.config.progress_strategy_ratio[0] if len(self.config.progress_strategy_ratio) > 0 else 0.0,84                ),85                (86                    DataGenStrat.FORWARD_PROGRESS,87                    self.config.progress_strategy_ratio[1] if len(self.config.progress_strategy_ratio) > 1 else 0.0,88                ),89                (90                    DataGenStrat.REVERSE_PROGRESS,91                    self.config.progress_strategy_ratio[2] if len(self.config.progress_strategy_ratio) > 2 else 0.0,92                ),93                (94                    DataGenStrat.REWIND,95                    self.config.progress_strategy_ratio[3] if len(self.config.progress_strategy_ratio) > 3 else 0.0,96                ),97            ]98 99            # Remove strategies with zero probability100            strategies = [(strat, prob) for strat, prob in strategies if prob > 0]101 102            max_attempts = 10  # Limit retry attempts to prevent infinite loops103            attempt = 0104 105            while processed_traj is None and attempt < max_attempts:106                attempt += 1107 108                # Check if we have any strategies left109                if not strategies:110                    return None111 112                # Rebalance probabilities based on remaining strategies113                total_prob = sum(prob for _, prob in strategies)114                if total_prob == 0:115                    return None116 117                # Normalize probabilities118                normalized_strategies = [(strat, prob / total_prob) for strat, prob in strategies]119 120                # Select strategy based on rebalanced probabilities121                prob = random.random()122                cumulative_prob = 0.0123                selected_strategy = None124 125                for strat, normalized_prob in normalized_strategies:126                    cumulative_prob += normalized_prob127                    if prob <= cumulative_prob:128                        selected_strategy = strat129                        break130 131                # Execute selected strategy132                result = self._execute_strategy(selected_strategy, traj)133                if result is not None:134                    processed_traj, subsample_strategy = result135                    strategy_used = selected_strategy136                else:137                    # Remove failed strategy and try again138                    strategies = [(strat, prob) for strat, prob in strategies if strat != selected_strategy]139                    continue140 141            # If we still don't have a sample after all attempts, return None142            if processed_traj is None:143                logger.trace(144                    f"[PROGRESS SAMPLER] Failed to generate progress sample after {max_attempts} attempts - all strategies exhausted"145                )146                return None147 148        progress_traj = self._get_traj_from_data(processed_traj, subsample_strategy=subsample_strategy)149 150        if progress_traj is None:151            return None152 153        # Handle special cases154        if strategy_used in [DataGenStrat.DIFFERENT_TASK, DataGenStrat.DIFFERENT_TASK_INSTRUCTION]:155            # We need to use the original task embeddings instead of the different task embeddings156            if self.config.load_embeddings and traj.get("embeddings_path"):157                progress_traj.text_embedding = load_embeddings_from_path(traj["embeddings_path"])["text_embedding"]158            progress_traj.lang_vector = traj["lang_vector"]159            progress_traj.task = traj["task"]160            progress_traj.target_progress = [0.0] * len(progress_traj.target_progress)161            if self.config.progress_loss_type.lower() == "discrete":162                progress_traj.target_progress = convert_continuous_to_discrete_bins(163                    progress_traj.target_progress, self.config.progress_discrete_bins164                )165            # Also set success labels to 0.0 (predict 0 success for different task trajectories)166            if progress_traj.success_label is not None:167                progress_traj.success_label = [0.0] * len(progress_traj.success_label)168 169        strategy_value = strategy_used.value if isinstance(strategy_used, DataGenStrat) else strategy_used170        sample = ProgressSample(171            trajectory=progress_traj,172            sample_type="progress",173            data_gen_strategy=strategy_value,174        )175        sample.resample_attempts = attempt176        return sample177