robometer/rewardeval_ui

1from typing import Dict, List, Any, Optional2from itertools import cycle3 4import numpy as np5from collections import defaultdict6from rfm.data.dataset_types import ProgressSample7from rfm.data.samplers.base import RFMBaseSampler8from rfm.utils.logger import get_logger9 10logger = get_logger()11 12 13class ProgressPolicyRankingSampler(RFMBaseSampler):14    """Dataset that generates progress samples for policy ranking by selecting N trajectories per quality label for tasks with multiple quality labels."""15 16    def __init__(17        self,18        num_examples_per_quality_pr: int = 5,19        num_partial_successes: Optional[int] = None,20        frame_step: int = 1,21        use_frame_steps: bool = True,22        max_tasks: Optional[int] = None,23        **kwargs,24    ):25        super().__init__(**kwargs)26 27        self.num_examples_per_quality_pr = num_examples_per_quality_pr28        self.num_partial_successes = num_partial_successes29        self.frame_step = frame_step30        self.use_frame_steps = use_frame_steps31        self.max_tasks = max_tasks32        logger.info(f"ProgressPolicyRankingSampler initialized with {len(self.robot_trajectories)} trajectories")33 34        self.sample_indices = self._generate_all_sample_indices()35 36        logger.info(f"Generated {len(self.sample_indices)} sample indices")37 38    def _generate_all_sample_indices(self) -> List[Dict[str, Any]]:39        """Generate sample indices by selecting tasks with multiple quality labels/partial_success values and sampling N trajectories per group.40 41        For non-RoboArena: Groups by task and quality_label.42        For RoboArena: Groups by task and partial_success values.43 44        If use_frame_steps=True, generates subsequence samples like reward_alignment (0:frame_step, 0:2*frame_step, etc.).45        If use_frame_steps=False, generates one sample per trajectory (whole trajectory).46        """47 48        # Check if this is RoboArena (has partial_success)49        is_roboarena = False50        if self.robot_trajectories:51            first_traj = self.dataset[self.robot_trajectories[0]]52            is_roboarena = first_traj.get("partial_success") is not None53 54        # Group trajectories by task and grouping key (quality_label or partial_success)55        task_to_key_to_trajs = defaultdict(lambda: defaultdict(list))56 57        for traj_idx in self.robot_trajectories:58            traj = self.dataset[traj_idx]59            task = traj["task"]60 61            if is_roboarena:62                # RoboArena: Use rounded partial_success as key to group similar values63                partial_success_val = traj.get("partial_success")64                if partial_success_val is not None:65                    partial_success = round(float(partial_success_val), 2)66                    task_to_key_to_trajs[task][partial_success].append(traj_idx)67            else:68                # Non-RoboArena: Use quality_label69                quality = traj["quality_label"]70                task_to_key_to_trajs[task][quality].append(traj_idx)71 72        # Filter to tasks that have multiple grouping values73        tasks_with_multiple_values = {74            task: key_to_trajs for task, key_to_trajs in task_to_key_to_trajs.items() if len(key_to_trajs) > 175        }76 77        dataset_type_str = "partial_success values" if is_roboarena else "quality labels"78        logger.info(f"Found {len(tasks_with_multiple_values)} tasks with multiple {dataset_type_str}")79 80        # Limit number of tasks if max_tasks is specified81        if self.max_tasks is not None and self.max_tasks > 0:82            # Convert to list, shuffle, and take first max_tasks83            # Sort tasks first to ensure deterministic ordering before shuffling84            tasks_list = sorted(tasks_with_multiple_values.items())85            self._local_random.shuffle(tasks_list)86            tasks_with_multiple_values = dict(tasks_list[: self.max_tasks])87            logger.info(f"Limited to {len(tasks_with_multiple_values)} tasks (max_tasks={self.max_tasks})")88 89        # Sample trajectories for each task90        sample_indices = []91        all_sampled_traj_indices = []92        # Sort tasks to ensure deterministic processing order93        for task, key_to_trajs in sorted(tasks_with_multiple_values.items()):94            if is_roboarena:95                # RoboArena: Use num_partial_successes for circular sampling96                num_to_sample_total = self.num_partial_successes97 98                # Build lists of available indices per partial_success (sorted for deterministic sampling)99                available_lists = []100                for partial_success in sorted(key_to_trajs.keys()):101                    traj_indices = sorted(key_to_trajs[partial_success])102                    if traj_indices:103                        available_lists.append(traj_indices)104 105                # Circular sampling: cycle through partial_success groups until we reach max106                sampled_traj_indices = []107                for available_indices in cycle(available_lists):108                    if len(sampled_traj_indices) >= num_to_sample_total:109                        break110                    if not available_indices:111                        # If all lists are empty, stop112                        if all(not lst for lst in available_lists):113                            break114                        continue115 116                    # Sample one trajectory from this group117                    sampled_idx = self._local_random.choice(available_indices)118                    sampled_traj_indices.append(sampled_idx)119                    # Remove the sampled index from this list120                    available_indices.remove(sampled_idx)121 122                # Generate samples for all sampled trajectories123                for traj_idx in sampled_traj_indices:124                    traj = self.dataset[traj_idx]125                    sample_indices.extend(self._generate_indices_for_trajectory(traj_idx, traj))126                    all_sampled_traj_indices.append(traj_idx)127            else:128                # Non-RoboArena: Sample N trajectories per quality label129                # Sort quality labels to ensure deterministic order130                for quality in sorted(key_to_trajs.keys()):131                    traj_indices = key_to_trajs[quality]132                    # Sort trajectory indices to ensure deterministic sampling133                    traj_indices = sorted(traj_indices)134                    # Sample up to num_examples_per_quality_pr trajectories for this quality label135                    num_to_sample = min(self.num_examples_per_quality_pr, len(traj_indices))136                    sampled_traj_indices = self._local_random.sample(traj_indices, num_to_sample)137                    for traj_idx in sampled_traj_indices:138                        traj = self.dataset[traj_idx]139                        sample_indices.extend(self._generate_indices_for_trajectory(traj_idx, traj))140                        all_sampled_traj_indices.append(traj_idx)141 142        logger.info(f"Sampled {len(sample_indices)} samples across {len(tasks_with_multiple_values)} tasks")143        logger.info(f"Sampled trajectory indices: {all_sampled_traj_indices}")144 145        return sample_indices146 147    def _generate_indices_for_trajectory(self, traj_idx: int, traj: Dict[str, Any]) -> List[Dict[str, Any]]:148        """Generate sample indices for a single trajectory.149 150        Args:151            traj_idx: Index of the trajectory in the dataset152            traj: Trajectory dictionary153 154        Returns:155            List of sample index dictionaries156        """157        num_frames = traj["num_frames"]158        indices = []159 160        if self.use_frame_steps:161            # Generate subsequence indices like reward_alignment: 0:frame_step, 0:2*frame_step, etc.162            for end_idx in range(self.frame_step, num_frames + 1, self.frame_step):163                frame_indices = list(range(end_idx))164                indices.append({165                    "traj_idx": traj_idx,166                    "frame_indices": frame_indices,167                    "num_frames": num_frames,168                    "video_path": traj["frames"],169                    "id": traj["id"],170                    "use_frame_steps": True,171                })172        else:173            # Generate one sample per trajectory (whole trajectory)174            indices.append({175                "traj_idx": traj_idx,176                "video_path": traj["frames"],177                "id": traj["id"],178                "use_frame_steps": False,179            })180 181        return indices182 183    def _generate_sample_from_indices(self, sample_idx_info: dict) -> ProgressSample:184        """Generate a single progress sample from trajectory index."""185        traj_idx = sample_idx_info["traj_idx"]186        use_frame_steps = sample_idx_info.get("use_frame_steps", True)187 188        traj = self.dataset[traj_idx]189 190        if use_frame_steps:191            # Frame steps mode: create subsequence like reward_alignment192            frame_indices = sample_idx_info["frame_indices"]193            num_frames = sample_idx_info["num_frames"]194 195            metadata = {196                "quality_label": traj["quality_label"],197                "data_source": traj["data_source"],198                "task": traj["task"],199                "id": traj["id"],200                "video_path": sample_idx_info["video_path"],201                "frame_step": frame_indices[-1] if frame_indices else 0,202            }203 204            trajectory = self._get_traj_from_data(205                traj=traj,206                frame_indices=frame_indices,207                metadata=metadata,208            )209        else:210            # Whole trajectory mode211            metadata = {212                "quality_label": traj["quality_label"],213                "data_source": traj["data_source"],214                "task": traj["task"],215                "id": traj["id"],216                "video_path": sample_idx_info["video_path"],217            }218 219            trajectory = self._get_traj_from_data(220                traj=traj,221                metadata=metadata,222            )223 224        sample = ProgressSample(trajectory=trajectory)225        return sample226 227    def __len__(self):228        return len(self.sample_indices)229 230    def __getitem__(self, idx):231        return self._generate_sample_from_indices(self.sample_indices[idx])232