robocandywrapper 0.2.2__tar.gz → 0.2.4__tar.gz

{robocandywrapper-0.2.2/robocandywrapper.egg-info → robocandywrapper-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robocandywrapper
-Version: 0.2.2
+Version: 0.2.4
 Summary: Sweet wrappers for extending and remixing LeRobot Datasets
 Author: RoboCandyWrapper Contributors
 License: MIT License

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "robocandywrapper"
-version = "0.2.2"
+version = "0.2.4"
 description = "Sweet wrappers for extending and remixing LeRobot Datasets"
 readme = "README.md"
 requires-python = ">=3.10"

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4}/robocandywrapper/__init__.py

@@ -17,7 +17,7 @@ from robocandywrapper.constants import (
     EPISODE_OUTCOME_PLUGIN_NAME,
 )
 
-__version__ = "0.2.1"
+__version__ = "0.2.4"
 
 __all__ = [
     "DatasetPlugin",

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4}/robocandywrapper/dataformats/lerobot_21/dataset.py

@@ -93,6 +93,7 @@ class LeRobot21DatasetMetadata:
         try:
             if force_cache_sync:
                 raise FileNotFoundError
+            self.pull_from_repo()
             self.load_metadata()
         except (FileNotFoundError, NotADirectoryError):
             if is_valid_version(self.revision):
@@ -728,7 +729,11 @@ class LeRobot21Dataset(torch.utils.data.Dataset):
         item = {}
         for vid_key, query_ts in query_timestamps.items():
             video_path = self.root / self.meta.get_video_file_path(ep_idx, vid_key)
-            frames = decode_video_frames(video_path, query_ts, self.tolerance_s, self.video_backend)
+            try:
+                frames = decode_video_frames(video_path, query_ts, self.tolerance_s, self.video_backend)
+            except Exception as e:
+                # fall back to trying to decode with pyav
+                frames = decode_video_frames(video_path, query_ts, self.tolerance_s, "pyav")
             item[vid_key] = frames.squeeze(0)
 
         return item
@@ -768,6 +773,11 @@ class LeRobot21Dataset(torch.utils.data.Dataset):
         task_idx = item["task_index"].item()
         item["task"] = self.meta.tasks[task_idx]
 
+        # Hack - add gripper position to end
+        # only applies to a specific dataset
+        # if "observation.eef_6d_pose" in item and item["observation.eef_6d_pose"].shape[0] == 6:
+        #     item["observation.eef_6d_pose"] = torch.cat([item["observation.eef_6d_pose"], item["observation.state"][-1:]], dim=0)
+
         return item
 
     def __repr__(self):

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4}/robocandywrapper/factory.py

@@ -59,7 +59,7 @@ def resolve_delta_timestamps(
     for key in ds_meta.features:
         if key == REWARD and reward_indices is not None:
             delta_timestamps[key] = _indices_to_times(reward_indices, ds_meta.fps)
-        if key == ACTION and action_indices is not None:
+        if key.startswith(ACTION) and action_indices is not None:
             delta_timestamps[key] = _indices_to_times(action_indices, ds_meta.fps)
         if key.startswith("observation.") and observation_indices is not None:
             delta_timestamps[key] = _indices_to_times(observation_indices, ds_meta.fps)
@@ -76,6 +76,7 @@ def _create_datasets(
     observation_delta_indices: Optional[List] = None,
     reward_delta_indices: Optional[List] = None,
     use_imagenet_stats: bool = True,
+    load_videos: bool = True,
 ) -> List[LeRobotDataset | LeRobot21Dataset]:
     """Private helper to create dataset instances from a list of repo IDs.
     
@@ -92,6 +93,8 @@ def _create_datasets(
         observation_delta_indices: Frame indices for observations.
         reward_delta_indices: Frame indices for rewards.
         use_imagenet_stats: Whether to apply ImageNet normalization stats.
+        load_videos: Whether to download and load video files (default: True).
+            Set to False to skip video downloads when not needed.
     
     Returns:
         List of dataset instances.
@@ -170,6 +173,7 @@ def _create_datasets(
             image_transforms=None,  # Will be applied by WrappedRobotDataset
             revision=revision,
             video_backend=video_backend,
+            download_videos=load_videos,
         )
         
         # Apply ImageNet stats if needed
@@ -189,12 +193,18 @@ def _create_datasets(
 def make_dataset(
     cfg: TrainPipelineConfig,
     plugins: Optional[list[DatasetPlugin]] = None,
+    key_rename_map: Optional[dict[str, str]] = None,
+    load_videos: bool = True,
 ) -> WrappedRobotDataset:
     """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
 
     Args:
         cfg (TrainPipelineConfig): A TrainPipelineConfig config which contains a DatasetConfig and a PreTrainedConfig.
         plugins (Optional[list[DatasetPlugin]]): Optional list of plugins to attach to the dataset(s).
+        key_rename_map (Optional[dict[str, str]]): Optional mapping from source keys to target keys
+            for unifying datasets with different naming conventions. Example: {"action.pos": "action"}
+        load_videos (bool): Whether to download and load video files (default: True).
+            Set to False to skip video downloads when not needed.
 
     Returns:
         WrappedRobotDataset: A wrapped dataset with plugin support.
@@ -221,6 +231,7 @@ def make_dataset(
         observation_delta_indices=cfg.policy.observation_delta_indices,
         reward_delta_indices=cfg.policy.reward_delta_indices,
         use_imagenet_stats=cfg.dataset.use_imagenet_stats,
+        load_videos=load_videos,
     )
     
     # Wrap in WrappedRobotDataset with plugins
@@ -228,6 +239,7 @@ def make_dataset(
         datasets=datasets,
         plugins=plugins,
         image_transforms=image_transforms,
+        key_rename_map=key_rename_map,
     )
     
     return wrapped_dataset
@@ -243,6 +255,8 @@ def make_dataset_without_config(
     revision: str | None = None,
     use_imagenet_stats: bool = True,
     plugins: Optional[list[DatasetPlugin]] = None,
+    key_rename_map: Optional[dict[str, str]] = None,
+    load_videos: bool = True,
 ) -> WrappedRobotDataset:
     """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
 
@@ -259,6 +273,10 @@ def make_dataset_without_config(
         revision (str, optional): Dataset revision
         use_imagenet_stats (bool): Whether to use ImageNet normalization stats (default: True)
         plugins (Optional[list[DatasetPlugin]]): Optional list of plugins to attach to the dataset(s)
+        key_rename_map (Optional[dict[str, str]]): Optional mapping from source keys to target keys
+            for unifying datasets with different naming conventions. Example: {"action.pos": "action"}
+        load_videos (bool): Whether to download and load video files (default: True).
+            Set to False to skip video downloads when not needed.
 
     Returns:
         WrappedRobotDataset: A wrapped dataset with plugin support.
@@ -283,12 +301,14 @@ def make_dataset_without_config(
         action_delta_indices=action_delta_indices,
         observation_delta_indices=observation_delta_indices,
         use_imagenet_stats=use_imagenet_stats,
+        load_videos=load_videos,
     )
     
     # Wrap in WrappedRobotDataset with plugins
     wrapped_dataset = WrappedRobotDataset(
         datasets=datasets,
         plugins=plugins,
+        key_rename_map=key_rename_map,
     )
     
     return wrapped_dataset

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4}/robocandywrapper/metadata_view.py

@@ -102,6 +102,7 @@ class WrappedRobotDatasetMetadataView:
         datasets: list,
         plugin_instances: list[list],
         dataset_weights: Optional[dict[str, float]] = None,
+        dataset_renames: Optional[list[dict[str, str]]] = None,
     ):
         """
         Initialize metadata view.
@@ -110,15 +111,32 @@ class WrappedRobotDatasetMetadataView:
             datasets: List of LeRobotDataset instances
             plugin_instances: List of plugin instances for each dataset
             dataset_weights: Optional weights for each dataset (for weighted stats)
+            dataset_renames: Optional list of rename dicts for each dataset,
+                mapping source_key -> target_key. Used to unify keys across
+                datasets with different naming conventions.
         """
         self._datasets = datasets
         self._plugin_instances = plugin_instances
         self._dataset_weights = dataset_weights or {}
+        self._dataset_renames = dataset_renames or [{} for _ in datasets]
         
         # Cache computed properties
         self._features = None
         self._stats = None
     
+    def _get_renamed_features(self, dataset_idx: int) -> dict[str, dict]:
+        """Get features from a dataset with key renames applied."""
+        dataset = self._datasets[dataset_idx]
+        renames = self._dataset_renames[dataset_idx]
+        
+        renamed_features = {}
+        for key, value in dataset.meta.features.items():
+            # Apply rename if applicable
+            effective_key = renames.get(key, key)
+            renamed_features[effective_key] = value
+        
+        return renamed_features
+
     @property
     def features(self) -> dict[str, dict]:
         """
@@ -127,6 +145,9 @@ class WrappedRobotDatasetMetadataView:
         Returns intersection of:
         1. Features from all datasets (taking intersection, not union)
         2. Features provided by plugins (added to intersection)
+        
+        Key renames are applied before computing the intersection, allowing
+        datasets with different naming conventions to be unified.
         """
         if self._features is not None:
             return self._features
@@ -135,12 +156,13 @@ class WrappedRobotDatasetMetadataView:
         if not self._datasets:
             all_features = {}
         else:
-            # Start with all features from first dataset
-            all_features = dict(self._datasets[0].meta.features)
+            # Start with all features from first dataset (with renames applied)
+            all_features = self._get_renamed_features(0)
             
             # Intersect with features from other datasets
-            for dataset in self._datasets[1:]:
-                dataset_feature_keys = set(dataset.meta.features.keys())
+            for i in range(1, len(self._datasets)):
+                dataset_features = self._get_renamed_features(i)
+                dataset_feature_keys = set(dataset_features.keys())
                 all_feature_keys = set(all_features.keys())
                 
                 # Keep only features that exist in both
@@ -166,6 +188,19 @@ class WrappedRobotDatasetMetadataView:
         self._features = all_features
         return self._features
     
+    def _get_renamed_stats(self, dataset_idx: int) -> dict[str, dict]:
+        """Get stats from a dataset with key renames applied."""
+        dataset = self._datasets[dataset_idx]
+        renames = self._dataset_renames[dataset_idx]
+        
+        renamed_stats = {}
+        for key, value in dataset.meta.stats.items():
+            # Apply rename if applicable
+            effective_key = renames.get(key, key)
+            renamed_stats[effective_key] = value
+        
+        return renamed_stats
+
     @property
     def stats(self) -> dict:
         """
@@ -174,12 +209,17 @@ class WrappedRobotDatasetMetadataView:
         If dataset_weights are provided, stats are computed as a weighted
         average based on effective dataset sizes (size * weight).
         Uses the correct statistical formula for combining variances.
+        
+        Key renames are applied before aggregation, so different source keys
+        (e.g., "action.pos" and "trajectory") that map to the same target key
+        (e.g., "action") will have their stats combined as if they were the
+        same key across all datasets.
         """
         if self._stats is not None:
             return self._stats
         
-        # Collect stats and weights for each dataset
-        stats_list = [dataset.meta.stats for dataset in self._datasets]
+        # Collect stats (with renames applied) and weights for each dataset
+        stats_list = [self._get_renamed_stats(i) for i in range(len(self._datasets))]
         
         # Get weight multiplier for each dataset
         weights = []

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4}/robocandywrapper/utils.py

@@ -10,11 +10,7 @@ from typing import Optional
 from glob import glob
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from lerobot.utils.constants import PRETRAINED_MODEL_DIR
-from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
-from lerobot.configs.types import FeatureType
-from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.datasets.utils import dataset_to_policy_features
 from termcolor import colored
 
 

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4}/robocandywrapper/wrapper.py

@@ -23,6 +23,7 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
         warn_on_key_conflicts: bool = True,
         error_on_key_conflicts: bool = True,
         dataset_weights: Optional[dict[str, float]] = None,
+        key_rename_map: Optional[dict[str, str]] = None,
         **kwargs
     ):
         """
@@ -35,6 +36,14 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
             warn_on_key_conflicts: Warn when plugins have overlapping keys (if not raising errors)
             error_on_key_conflicts: Raise error on key conflicts (default: True)
             dataset_weights: Optional weights for computing weighted stats (e.g., {"dataset_id": 2.0})
+            key_rename_map: Optional mapping from source keys to target keys for unifying
+                datasets with different naming conventions. Keys are renamed before the
+                intersection logic runs, allowing datasets with different key names to be
+                combined. Example: {"action.pos": "action", "trajectory": "action"}
+                
+                Note: When a key is renamed, any corresponding "_is_pad" key (added by
+                LeRobot when using delta_timestamps) is automatically renamed as well.
+                E.g., "action.pos" -> "action" also renames "action.pos_is_pad" -> "action_is_pad".
         """
         super().__init__()
         
@@ -64,6 +73,10 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
             self._cumulative_lengths.append(self._cumulative_lengths[-1] + length)
         self._total_length = self._cumulative_lengths[-1]
         
+        # Key rename mapping: unify differently-named keys across datasets
+        self.key_rename_map = key_rename_map or {}
+        self._dataset_renames = self._compute_dataset_renames()
+        
         # Plugin management: one plugin class, many instances (one per dataset)
         self._plugins: list[DatasetPlugin] = plugins or []
         self._plugin_instances: list[list[PluginInstance]] = []
@@ -84,6 +97,7 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
             datasets=self._datasets,
             plugin_instances=self._plugin_instances,
             dataset_weights=dataset_weights,
+            dataset_renames=self._dataset_renames,
         )
 
         # ** MATCHING LeRobot MULTI-DATASET API DESIGN **
@@ -91,41 +105,51 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
         # Disable any data keys that are not common across all of the datasets. Note: we may relax this
         # restriction in future iterations of this class. For now, this is necessary at least for being able
         # to use PyTorch's default DataLoader collate function.
+        #
+        # Key rename mapping is applied first (conceptually), so intersection is computed on
+        # "effective" features (post-rename). This allows datasets with different key names to be
+        # unified before the intersection check.
         self.disabled_features = set()
-        intersection_features = set(self._datasets[0].features)
-        for ds in self._datasets:
-            intersection_features.intersection_update(ds.features)
+        intersection_features = self._get_effective_features(0)
+        for i in range(len(self._datasets)):
+            intersection_features.intersection_update(self._get_effective_features(i))
         if len(intersection_features) == 0:
             raise RuntimeError(
                 "Multiple datasets were provided but they had no keys common to all of them. "
                 "The multi-dataset functionality currently only keeps common keys."
             )
-        for repo_id, ds in zip(self.repo_ids, self._datasets, strict=True):
-            extra_keys = set(ds.features).difference(intersection_features)
+        for i, repo_id in enumerate(self.repo_ids):
+            effective_keys = self._get_effective_features(i)
+            extra_keys = effective_keys.difference(intersection_features)
             if extra_keys:
                 logging.warning(
-                f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the "
-                "other datasets."
+                    f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the "
+                    "other datasets."
                 )
             self.disabled_features.update(extra_keys)
 
         # Validate that common features have compatible shapes
+        # Note: We need to look up the original key name for renamed keys
         for key in intersection_features:
             shapes = []
-            for ds in self._datasets:
-                if key in ds.meta.features:
-                    feature_shape = ds.meta.features[key].get('shape', [])
+            shape_details = []
+            for i, ds in enumerate(self._datasets):
+                # Find the original key (may be renamed)
+                renames = self._dataset_renames[i]
+                reverse_renames = {v: k for k, v in renames.items()}
+                original_key = reverse_renames.get(key, key)
+                
+                if original_key in ds.meta.features:
+                    feature_shape = ds.meta.features[original_key].get('shape', [])
                     shapes.append(tuple(feature_shape))
+                    if original_key != key:
+                        shape_details.append(f"{ds.repo_id}: {feature_shape} (from '{original_key}')")
+                    else:
+                        shape_details.append(f"{ds.repo_id}: {feature_shape}")
             
             # Check if all shapes are the same
             unique_shapes = set(shapes)
             if len(unique_shapes) > 1:
-                shape_details = []
-                for ds in self._datasets:
-                    if key in ds.meta.features:
-                        shape = ds.meta.features[key].get('shape', [])
-                        shape_details.append(f"{ds.repo_id}: {shape}")
-                
                 raise ValueError(
                     f"Incompatible shapes for feature '{key}' across datasets:\n" +
                     "\n".join(f"  - {detail}" for detail in shape_details) +
@@ -208,8 +232,11 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
         plugin_only_features = {}
         for key, value in self._meta.features.items():
             if key not in base_features:
-                plugin_only_features[key] = PolicyFeature(type=FeatureType.STATE, shape=value['shape'])
-        
+                if 'action' in key:
+                    plugin_only_features[key] = PolicyFeature(type=FeatureType.ACTION, shape=value['shape'])
+                else:
+                    plugin_only_features[key] = PolicyFeature(type=FeatureType.STATE, shape=value['shape'])
+
         return plugin_only_features
     
     @property
@@ -296,6 +323,77 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
         # Also update the cached stats property
         self.stats = self._meta.stats
 
+    def _compute_dataset_renames(self) -> list[dict[str, str]]:
+        """
+        Pre-compute which key renames apply to each dataset.
+        
+        For each dataset, determines which source keys from key_rename_map exist
+        and can be renamed (i.e., target key doesn't already exist).
+        
+        Also automatically handles derived _is_pad keys that LeRobot adds when
+        delta_timestamps are used. For example, if renaming "action.pos" -> "action",
+        this will also rename "action.pos_is_pad" -> "action_is_pad".
+        
+        Returns:
+            List of dicts mapping source_key -> target_key for each dataset
+        """
+        dataset_renames = []
+        for dataset in self._datasets:
+            ds_renames = {}
+            ds_keys = set(dataset.features)
+            
+            for source, target in self.key_rename_map.items():
+                if source in ds_keys:
+                    if target in ds_keys:
+                        # Target already exists in this dataset - skip rename to avoid conflict
+                        logging.warning(
+                            f"Skipping rename '{source}' -> '{target}' for {dataset.repo_id}: "
+                            f"target key already exists in dataset"
+                        )
+                    else:
+                        ds_renames[source] = target
+                        
+                        # Also handle the _is_pad suffix that LeRobot adds for delta_timestamps
+                        # These keys are dynamically added during __getitem__ and may not be in
+                        # dataset.features, but we still want to rename them consistently
+                        is_pad_source = f"{source}_is_pad"
+                        is_pad_target = f"{target}_is_pad"
+                        
+                        # Check for conflicts on the _is_pad key as well
+                        if is_pad_target in ds_keys:
+                            logging.warning(
+                                f"Skipping derived rename '{is_pad_source}' -> '{is_pad_target}' "
+                                f"for {dataset.repo_id}: target key already exists in dataset"
+                            )
+                        else:
+                            ds_renames[is_pad_source] = is_pad_target
+            
+            dataset_renames.append(ds_renames)
+        
+        return dataset_renames
+    
+    def _get_effective_features(self, dataset_idx: int) -> set[str]:
+        """
+        Get the effective feature keys for a dataset after applying renames.
+        
+        Args:
+            dataset_idx: Index of the dataset
+            
+        Returns:
+            Set of feature keys that would exist after renaming
+        """
+        ds = self._datasets[dataset_idx]
+        renames = self._dataset_renames[dataset_idx]
+        
+        effective = set()
+        for key in ds.features:
+            if key in renames:
+                effective.add(renames[key])
+            else:
+                effective.add(key)
+        
+        return effective
+
     def _validate_plugin_keys(self):
         """
         Check for key conflicts between plugins.
@@ -483,7 +581,14 @@ class WrappedRobotDataset(torch.utils.data.Dataset):
         # Add dataset index
         item["dataset_index"] = torch.tensor(dataset_idx)
         
-        # Remove disabled features
+        # Apply key renaming for this dataset (before filtering disabled features)
+        # This unifies differently-named keys across datasets
+        renames = self._dataset_renames[dataset_idx]
+        for source, target in renames.items():
+            if source in item:
+                item[target] = item.pop(source)
+        
+        # Remove disabled features (now operates on effective/renamed key names)
         for data_key in self.disabled_features:
             if data_key in item:
                 del item[data_key]

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4/robocandywrapper.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robocandywrapper
-Version: 0.2.2
+Version: 0.2.4
 Summary: Sweet wrappers for extending and remixing LeRobot Datasets
 Author: RoboCandyWrapper Contributors
 License: MIT License

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4}/robocandywrapper.egg-info/SOURCES.txt

@@ -26,4 +26,5 @@ robocandywrapper/samplers/__init__.py
 robocandywrapper/samplers/config.py
 robocandywrapper/samplers/factory.py
 robocandywrapper/samplers/weighted.py
-tests/test_dataset_weights_integration.py
+tests/test_dataset_weights_integration.py
+tests/test_key_rename_stats.py

{robocandywrapper-0.2.2 → robocandywrapper-0.2.4}/setup.py

@@ -9,7 +9,7 @@ long_description = readme_file.read_text(encoding="utf-8") if readme_file.exists
 
 setup(
     name="robocandywrapper",
-    version="0.2.1",
+    version="0.2.4",
     description="Sweet wrappers for extending and remixing LeRobot Datasets",
     long_description=long_description,
     long_description_content_type="text/markdown",

robocandywrapper-0.2.4/tests/test_key_rename_stats.py

@@ -0,0 +1,394 @@
+"""
+Test for key_rename_map functionality in stats aggregation.
+
+Tests that when datasets have differently-named keys that map to the same 
+target key, their stats are properly combined as if they were the same key.
+"""
+
+import sys
+from pathlib import Path
+
+# Add the local package to path before imports
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+import numpy as np
+from robocandywrapper.wrapper import WrappedRobotDataset
+
+
+class MockLeRobotDataset:
+    """Mock dataset for testing."""
+    
+    def __init__(self, repo_id, fps, features, num_frames, stats=None):
+        self.repo_id = repo_id
+        self._fps = fps
+        self._features = features
+        self._num_frames = num_frames
+        
+        # Create mock metadata
+        self.meta = MockMetadata(repo_id, fps, features, stats)
+        self.hf_features = features
+        self.features = features
+    
+    def __len__(self):
+        return self._num_frames
+    
+    def __getitem__(self, idx):
+        # Minimal mock for dataset access
+        return {"action": np.array([0.0, 0.0])}
+
+
+class MockMetadata:
+    """Mock metadata object."""
+    
+    def __init__(self, repo_id, fps, features, stats=None):
+        self.repo_id = repo_id
+        self._fps = fps
+        self._features = features
+        self.info = {"fps": fps}
+        
+        # Use provided stats or default
+        if stats is None:
+            self.stats = {
+                "action": {
+                    "mean": np.array([0.0, 0.0]),
+                    "std": np.array([1.0, 1.0]),
+                    "min": np.array([-1.0, -1.0]),
+                    "max": np.array([1.0, 1.0]),
+                    "count": np.array([1000]),
+                }
+            }
+        else:
+            self.stats = stats
+            
+        self.camera_keys = [k for k in features if "image" in k or "video" in k]
+        self.image_keys = [k for k in features if "image" in k]
+        self.video_keys = [k for k in features if "video" in k]
+    
+    @property
+    def fps(self):
+        return self._fps
+    
+    @property
+    def features(self):
+        return self._features
+
+
+def test_key_rename_stats_aggregation():
+    """
+    Test that keys are properly renamed in stats aggregation.
+    
+    Scenario:
+    - Dataset 1 has "action.pos" key with certain stats
+    - Dataset 2 has "trajectory" key with certain stats
+    - key_rename_map maps both to "action"
+    - Result should have "action" stats that combine both sources
+    """
+    print("\n" + "="*60)
+    print("Test: Key Rename Stats Aggregation")
+    print("="*60)
+    
+    # Dataset 1: has "action.pos" key
+    stats1 = {
+        "action.pos": {
+            "mean": np.array([1.0, 2.0]),
+            "std": np.array([0.5, 0.5]),
+            "min": np.array([-1.0, -1.0]),
+            "max": np.array([3.0, 4.0]),
+            "count": np.array([1000]),  # 1000 samples
+        }
+    }
+    
+    # Dataset 2: has "trajectory" key  
+    stats2 = {
+        "trajectory": {
+            "mean": np.array([5.0, 6.0]),
+            "std": np.array([1.0, 1.0]),
+            "min": np.array([0.0, 0.0]),
+            "max": np.array([10.0, 12.0]),
+            "count": np.array([1000]),  # Same count for simpler math
+        }
+    }
+    
+    dataset1 = MockLeRobotDataset(
+        repo_id="dataset_with_action_pos",
+        fps=20,
+        features={"action.pos": {"shape": [2]}},
+        num_frames=1000,
+        stats=stats1
+    )
+    
+    dataset2 = MockLeRobotDataset(
+        repo_id="dataset_with_trajectory",
+        fps=20,
+        features={"trajectory": {"shape": [2]}},
+        num_frames=1000,
+        stats=stats2
+    )
+    
+    # Create wrapped dataset with key rename map
+    key_rename_map = {
+        "action.pos": "action",
+        "trajectory": "action",
+    }
+    
+    print(f"\n1. Creating wrapped dataset with key_rename_map: {key_rename_map}")
+    
+    wrapped_dataset = WrappedRobotDataset(
+        datasets=[dataset1, dataset2],
+        plugins=None,
+        key_rename_map=key_rename_map,
+    )
+    
+    # Check that features were renamed
+    print("\n2. Checking features")
+    assert "action" in wrapped_dataset.meta.features, \
+        "Renamed 'action' key should be in features"
+    assert "action.pos" not in wrapped_dataset.meta.features, \
+        "Original 'action.pos' key should not be in features"
+    assert "trajectory" not in wrapped_dataset.meta.features, \
+        "Original 'trajectory' key should not be in features"
+    print("   ✅ Features correctly renamed")
+    
+    # Check that stats were combined
+    print("\n3. Checking stats aggregation")
+    combined_stats = wrapped_dataset.meta.stats
+    
+    assert "action" in combined_stats, \
+        "Combined stats should have 'action' key"
+    assert "action.pos" not in combined_stats, \
+        "Original 'action.pos' should not be in combined stats"
+    assert "trajectory" not in combined_stats, \
+        "Original 'trajectory' should not be in combined stats"
+    
+    # With equal counts (1000 each), the combined mean should be the average
+    # mean = (1000 * [1.0, 2.0] + 1000 * [5.0, 6.0]) / 2000 = [3.0, 4.0]
+    expected_mean = np.array([3.0, 4.0])
+    
+    np.testing.assert_allclose(
+        combined_stats["action"]["mean"],
+        expected_mean,
+        rtol=1e-5,
+        err_msg="Combined mean should be average of both datasets' means"
+    )
+    print(f"   Combined mean: {combined_stats['action']['mean']}")
+    print(f"   Expected mean: {expected_mean}")
+    print("   ✅ Stats correctly combined")
+    
+    # Check min/max
+    expected_min = np.array([-1.0, -1.0])  # min of both datasets
+    expected_max = np.array([10.0, 12.0])  # max of both datasets
+    
+    np.testing.assert_allclose(
+        combined_stats["action"]["min"],
+        expected_min,
+        rtol=1e-5,
+        err_msg="Combined min should be minimum across both datasets"
+    )
+    np.testing.assert_allclose(
+        combined_stats["action"]["max"],
+        expected_max,
+        rtol=1e-5,
+        err_msg="Combined max should be maximum across both datasets"
+    )
+    print(f"   Combined min: {combined_stats['action']['min']}, expected: {expected_min}")
+    print(f"   Combined max: {combined_stats['action']['max']}, expected: {expected_max}")
+    print("   ✅ Min/max correctly combined")
+    
+    # Check count
+    assert combined_stats["action"]["count"] == 2000, \
+        f"Combined count should be 2000, got {combined_stats['action']['count']}"
+    print(f"   Combined count: {combined_stats['action']['count']}")
+    print("   ✅ Count correctly combined")
+    
+    print("\n" + "="*60)
+    print("✅ KEY RENAME STATS TEST PASSED!")
+    print("="*60 + "\n")
+
+
+def test_key_rename_with_different_counts():
+    """
+    Test key rename with datasets having different sample counts.
+    
+    This ensures weighted aggregation works correctly with renamed keys.
+    """
+    print("\n" + "="*60)
+    print("Test: Key Rename Stats with Different Counts")
+    print("="*60)
+    
+    # Dataset 1: 1000 samples with "pos"
+    stats1 = {
+        "pos": {
+            "mean": np.array([0.0]),
+            "std": np.array([1.0]),
+            "min": np.array([-3.0]),
+            "max": np.array([3.0]),
+            "count": np.array([1000]),
+        }
+    }
+    
+    # Dataset 2: 3000 samples with "position"
+    stats2 = {
+        "position": {
+            "mean": np.array([4.0]),
+            "std": np.array([2.0]),
+            "min": np.array([-2.0]),
+            "max": np.array([10.0]),
+            "count": np.array([3000]),
+        }
+    }
+    
+    dataset1 = MockLeRobotDataset(
+        repo_id="ds1",
+        fps=20,
+        features={"pos": {"shape": [1]}},
+        num_frames=1000,
+        stats=stats1
+    )
+    
+    dataset2 = MockLeRobotDataset(
+        repo_id="ds2",
+        fps=20,
+        features={"position": {"shape": [1]}},
+        num_frames=3000,
+        stats=stats2
+    )
+    
+    key_rename_map = {
+        "pos": "state",
+        "position": "state",
+    }
+    
+    print(f"\n1. Dataset 1: 1000 samples, mean=0.0")
+    print(f"   Dataset 2: 3000 samples, mean=4.0")
+    print(f"   key_rename_map: {key_rename_map}")
+    
+    wrapped_dataset = WrappedRobotDataset(
+        datasets=[dataset1, dataset2],
+        plugins=None,
+        key_rename_map=key_rename_map,
+    )
+    
+    combined_stats = wrapped_dataset.meta.stats
+    
+    # Expected weighted mean: (1000 * 0.0 + 3000 * 4.0) / 4000 = 3.0
+    expected_mean = np.array([3.0])
+    
+    print(f"\n2. Combined stats for 'state':")
+    print(f"   Mean: {combined_stats['state']['mean']} (expected: {expected_mean})")
+    
+    np.testing.assert_allclose(
+        combined_stats["state"]["mean"],
+        expected_mean,
+        rtol=1e-5,
+        err_msg="Weighted mean should account for different counts"
+    )
+    
+    assert combined_stats["state"]["count"] == 4000, \
+        f"Total count should be 4000, got {combined_stats['state']['count']}"
+    print(f"   Count: {combined_stats['state']['count']} (expected: 4000)")
+    
+    print("\n" + "="*60)
+    print("✅ KEY RENAME WITH DIFFERENT COUNTS TEST PASSED!")
+    print("="*60 + "\n")
+
+
+def test_key_rename_partial_rename():
+    """
+    Test that keys that don't need renaming are preserved.
+    
+    Only keys in key_rename_map should be renamed; others should pass through.
+    """
+    print("\n" + "="*60)
+    print("Test: Partial Key Rename")
+    print("="*60)
+    
+    # Both datasets have "action" (no rename needed) but different secondary keys
+    stats1 = {
+        "action": {
+            "mean": np.array([1.0]),
+            "std": np.array([1.0]),
+            "min": np.array([0.0]),
+            "max": np.array([2.0]),
+            "count": np.array([1000]),
+        },
+        "observation.state": {
+            "mean": np.array([0.5]),
+            "std": np.array([0.1]),
+            "min": np.array([0.0]),
+            "max": np.array([1.0]),
+            "count": np.array([1000]),
+        }
+    }
+    
+    stats2 = {
+        "action": {
+            "mean": np.array([3.0]),
+            "std": np.array([1.0]),
+            "min": np.array([2.0]),
+            "max": np.array([4.0]),
+            "count": np.array([1000]),
+        },
+        "observation.state": {
+            "mean": np.array([0.5]),
+            "std": np.array([0.1]),
+            "min": np.array([0.0]),
+            "max": np.array([1.0]),
+            "count": np.array([1000]),
+        }
+    }
+    
+    dataset1 = MockLeRobotDataset(
+        repo_id="ds1",
+        fps=20,
+        features={"action": {"shape": [1]}, "observation.state": {"shape": [1]}},
+        num_frames=1000,
+        stats=stats1
+    )
+    
+    dataset2 = MockLeRobotDataset(
+        repo_id="ds2",
+        fps=20,
+        features={"action": {"shape": [1]}, "observation.state": {"shape": [1]}},
+        num_frames=1000,
+        stats=stats2
+    )
+    
+    # No key rename map - should work normally
+    print("\n1. Creating wrapped dataset without key_rename_map")
+    
+    wrapped_dataset = WrappedRobotDataset(
+        datasets=[dataset1, dataset2],
+        plugins=None,
+        key_rename_map=None,
+    )
+    
+    combined_stats = wrapped_dataset.meta.stats
+    
+    # "action" should be combined normally
+    assert "action" in combined_stats, "action key should be present"
+    expected_action_mean = np.array([2.0])  # (1.0 + 3.0) / 2
+    
+    np.testing.assert_allclose(
+        combined_stats["action"]["mean"],
+        expected_action_mean,
+        rtol=1e-5
+    )
+    print(f"   action mean: {combined_stats['action']['mean']} (expected: {expected_action_mean})")
+    
+    # "observation.state" should also be present  
+    assert "observation.state" in combined_stats, "observation.state should be present"
+    print(f"   observation.state mean: {combined_stats['observation.state']['mean']}")
+    
+    print("\n" + "="*60)
+    print("✅ PARTIAL KEY RENAME TEST PASSED!")
+    print("="*60 + "\n")
+
+
+if __name__ == "__main__":
+    test_key_rename_stats_aggregation()
+    test_key_rename_with_different_counts()
+    test_key_rename_partial_rename()
+    
+    print("\n" + "="*60)
+    print("ALL KEY RENAME STATS TESTS PASSED!")
+    print("="*60 + "\n")