@@ -31,6 +31,7 @@ import torch.utils
from huggingface_hub import HfApi , snapshot_download
from huggingface_hub . errors import RevisionNotFoundError
from collections import defaultdict
from lerobot . constants import HF_LEROBOT_HOME
from lerobot . datasets . compute_stats import aggregate_stats , compute_episode_stats
from lerobot . datasets . image_writer import AsyncImageWriter , write_image
@@ -81,7 +82,12 @@ from lerobot.datasets.video_utils import (
)
CODEBASE_VERSION = " v3.0 "
OBS_IMAGE = " observation.image "
OBS_IMAGE_2 = " observation.image_2 "
OBS_IMAGE_3 = " observation.image_3 "
OBS_STATE = " observation.state "
OBS_ENV_STATE = " observation.env_state "
ACTION = " action "
class LeRobotDatasetMetadata :
def __init__ (
@@ -1322,13 +1328,139 @@ class LeRobotDataset(torch.utils.data.Dataset):
obj . video_backend = video_backend if video_backend is not None else get_safe_default_codec ( )
return obj
ROBOT_TYPE_KEYS_MAPPING = {
" lerobot/stanford_hydra_dataset " : " static_single_arm " ,
" lerobot/iamlab_cmu_pickup_insert " : " static_single_arm " ,
" lerobot/berkeley_fanuc_manipulation " : " static_single_arm " ,
" lerobot/toto " : " static_single_arm " ,
" lerobot/roboturk " : " static_single_arm " ,
" lerobot/jaco_play " : " static_single_arm " ,
" lerobot/taco_play " : " static_single_arm_7statedim " ,
}
class MultiLeRobotDatasetMeta :
def __init__ (
self ,
datasets : list [ LeRobotDataset ] ,
repo_ids : list [ str ] ,
keys_to_max_dim : dict [ str , int ] ,
train_on_all_features : bool = False ,
) :
self . repo_ids = repo_ids
self . keys_to_max_dim = keys_to_max_dim
self . train_on_all_features = train_on_all_features
self . robot_types = [ ds . meta . info [ " robot_type " ] for ds in datasets ]
# assign robot_type if missing
for ds in datasets :
ds . meta . info [ " robot_type " ] = ROBOT_TYPE_KEYS_MAPPING . get ( ds . repo_id , ds . meta . info [ " robot_type " ] )
ds . robot_type = ds . meta . info [ " robot_type " ]
# step 1: compute disabled features
self . disabled_features = set ( )
if not self . train_on_all_features :
intersection = set ( datasets [ 0 ] . features )
for ds in datasets :
intersection . intersection_update ( ds . features )
if not intersection :
raise RuntimeError ( " No common features across datasets. " )
for repo_id , ds in zip ( repo_ids , datasets , strict = False ) :
extra = set ( ds . features ) - intersection
logging . warning ( f " Disabling { extra } for repo { repo_id } " )
self . disabled_features . update ( extra )
# step 2: build union_features excluding disabled
self . union_features = { }
for ds in datasets :
for k , v in ds . features . items ( ) :
if k not in self . disabled_features :
self . union_features [ k ] = v
# step 3: reshape feature schema
self . features = reshape_features_to_max_dim (
self . union_features , reshape_dim = - 1 , keys_to_max_dim = self . keys_to_max_dim
)
# step 4: aggregate stats
self . stats = aggregate_stats_per_robot_type ( datasets )
for robot_type_ , stats_ in self . stats . items ( ) :
for feat_key , feat_stats in stats_ . items ( ) :
if feat_key in [ ACTION , OBS_ENV_STATE , OBS_STATE ] :
for k , v in feat_stats . items ( ) :
pad_value = 0 if k in [ " min " , " mean " ] else 1
self . stats [ robot_type_ ] [ feat_key ] [ k ] = pad_tensor (
v ,
max_size = self . keys_to_max_dim . get ( feat_key , - 1 ) ,
pad_dim = - 1 ,
pad_value = pad_value ,
)
# step 5: episodes & tasks
self . episodes = { repo_id : ds . meta . episodes for repo_id , ds in zip ( repo_ids , datasets , strict = False ) }
self . tasks = { repo_id : ds . meta . tasks for repo_id , ds in zip ( repo_ids , datasets , strict = False ) }
self . info = { repo_id : ds . meta . info for repo_id , ds in zip ( repo_ids , datasets , strict = False ) }
class MultiLeRobotDatasetCleaner :
def __init__ (
self ,
datasets : list [ LeRobotDataset ] ,
repo_ids : list [ str ] ,
sampling_weights : list [ float ] ,
datasets_repo_ids : list [ str ] ,
min_fps : int = 1 ,
max_fps : int = 100 ,
) :
self . original_datasets = datasets
self . original_repo_ids = repo_ids
self . original_weights = sampling_weights
self . original_datasets_repo_ids = datasets_repo_ids
# step 1: remove datasets with invalid fps
# step 2: keep datasets with same features per robot type
consistent_datasets , keep_mask = keep_datasets_with_the_same_features_per_robot_type (
datasets
)
self . cleaned_datasets = consistent_datasets
self . keep_mask = keep_mask
self . cleaned_weights = [ sampling_weights [ i ] for i in range ( len ( datasets ) ) if keep_mask [ i ] ]
self . cleaned_repo_ids = [ repo_ids [ i ] for i in range ( len ( datasets ) ) if keep_mask [ i ] ]
self . cleaned_datasets_repo_ids = [
datasets_repo_ids [ i ] for i in range ( len ( datasets ) ) if keep_mask [ i ]
]
self . cumulative_sizes = np . array (
[ 0 ] + list ( torch . cumsum ( torch . tensor ( [ len ( d ) for d in consistent_datasets ] ) , dim = 0 ) )
)
self . cleaned_weights = np . array ( self . cleaned_weights , dtype = np . float32 )
# --- at the top of the file (same imports as before) ---
from collections import defaultdict
from typing import Callable
import copy
import numpy as np
import torch
import datasets
from pathlib import Path
# If you already have these in your codebase, reuse them
try :
from lerobot . common . constants import (
ACTION , OBS_ENV_STATE , OBS_STATE , OBS_IMAGE , OBS_IMAGE_2 , OBS_IMAGE_3
)
except Exception :
# Fallbacks if constants are already strings elsewhere
ACTION = " action "
OBS_ENV_STATE = " observation.env_state "
OBS_STATE = " observation.state "
OBS_IMAGE = " observation.image "
OBS_IMAGE_2 = " observation.image_2 "
OBS_IMAGE_3 = " observation.image_3 "
IGNORED_KEYS = [ " observation.effort " ]
class MultiLeRobotDataset ( torch . utils . data . Dataset ) :
""" A dataset consisting of multiple underlying `LeRobotDataset`s.
The underlying `LeRobotDataset`s are effectively concatenated, and this class adopts much of the API
structure of `LeRobotDataset`.
"""
# ... keep your existing docstring ...
def __init__ (
self ,
@@ -1336,99 +1468,253 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
root : str | Path | None = None ,
episodes : dict | None = None ,
image_transforms : Callable | None = None ,
delta_timestamps : dict [ str , list [ float ] ] | None = None ,
delta_timestamps : dict [ list [ float ] ] | None = None ,
tolerances_s : dict | None = None ,
download_videos : bool = True ,
video_backend : str | None = None ,
# --- NEW: simple add-ons ---
sampling_weights : list [ float ] | None = None ,
feature_keys_mapping : dict [ str , dict [ str , str ] ] | None = None ,
max_action_dim : int | None = None ,
max_state_dim : int | None = None ,
max_num_images : int | None = None ,
max_image_dim : int | None = None ,
train_on_all_features : bool = False ,
min_fps : int = 1 ,
max_fps : int = 100 ,
ignore_keys : list [ str ] | None = None , # exact or glob patterns
) :
super ( ) . __init__ ( )
self . repo_ids = repo_ids
self . root = Path ( root ) if root else HF_LEROBOT_HOME
self . tolerances_s = tolerances_s if tolerances_s else dict . fromkeys ( repo_ids , 0.0001 )
# Construct the underlying datasets passing everything but `transform` and `delta_timestamps` which
# are handled by this class.
self . _datasets = [
LeRobotDataset (
repo_id ,
root = self . root / repo_id ,
episodes = episodes [ repo_id ] if episodes else None ,
image_transforms = image_transforms ,
delta_timestamps = delta_timestamps ,
tolerance_s = self . tolerances_s [ repo_id ] ,
download_videos = download_videos ,
video_backend = video_backend ,
)
for repo_id in repo_ids
]
# 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.
self . disabled_features = set ( )
intersection_features = set ( self . _datasets [ 0 ] . features )
for ds in self . _datasets :
intersection_features . intersection_update ( ds . features )
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 )
logging . warning (
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 )
# --- NEW: store mapping and simple knobs ---
self . feature_keys_mapping : dict [ str , dict [ str , str ] ] = feature_keys_mapping or { }
self . train_on_all_features = train_on_all_features
self . max_action_dim = max_action_dim
self . max_state_dim = max_state_dim
self . max_image_dim = max_image_dim
self . max_num_images = max_num_images # (optional, we don’
self . _ignore_patterns = list ( ignore_keys or [ ] )
# Build underlying single datasets
_datasets = [ ]
datasets_repo_ids = [ ]
self . sampling_weights = [ ]
sampling_weights = sampling_weights if sampling_weights is not None else [ 1 ] * len ( repo_ids )
assert len ( sampling_weights ) == len ( repo_ids ) , (
" The number of sampling weights must match the number of datasets. "
f " Got { len ( sampling_weights ) } weights for { len ( repo_ids ) } datasets. "
)
for i , repo_id in enumerate ( repo_ids ) :
try :
_datasets . append (
LeRobotDataset (
repo_id ,
root = self . root / repo_id ,
episodes = episodes . get ( repo_id , None ) if episodes else None ,
image_transforms = image_transforms , # transforms applied inside single ds
delta_timestamps = delta_timestamps . get ( repo_id , None ) if delta_timestamps else None ,
tolerance_s = self . tolerances_s [ repo_id ] ,
download_videos = download_videos ,
video_backend = video_backend ,
)
)
datasets_repo_ids . append ( repo_id )
self . sampling_weights . append ( float ( sampling_weights [ i ] ) )
except Exception as e :
print ( f " Failed to load dataset: { repo_id } due to Exception: { e } " )
print (
f " Finish loading { len ( _datasets ) } datasets, with sampling weights: "
f " { self . sampling_weights } corresponding to: { datasets_repo_ids } "
)
# Bookkeeping for mapping & canonical image inventory
self . image_transforms = image_transforms
self . delta_timestamps = delta_timestamps
# TODO(rcadene, aliberts): We should not perform this aggregation for datasets
# with multiple robots of different ranges. Instead we should have one normalization
# per robot.
self . stats = aggregate_stats ( [ dataset . meta . stats for dataset in self . _datasets ] )
self . delta_timestamps = delta_timestamps . get ( repo_id , None ) if delta_timestamps else None
self . _datasets = _datasets
self . datasets_repo_ids = datasets_repo_ids
# --- NEW: compute “canonical image keys” (targets across all mappings) ---
self . _canonical_image_keys : set [ str ] = set ( )
self . _source_keys_per_repo : dict [ str , set [ str ] ] = { }
self . _target_keys_per_repo : dict [ str , set [ str ] ] = { }
for rid , mapping in self . feature_keys_mapping . items ( ) :
src_keys = set ( mapping . keys ( ) )
tgt_keys = set ( mapping . values ( ) )
self . _source_keys_per_repo [ rid ] = src_keys
self . _target_keys_per_repo [ rid ] = tgt_keys
# union of target names (we will ensure these exist at __getitem__)
self . _canonical_image_keys | = {
k for k in tgt_keys if self . _is_image_key_like ( k )
}
# If user didn’
if not self . _canonical_image_keys and self . train_on_all_features :
# discover all image-like keys from raw features
for ds in self . _datasets :
for k , v in ds . hf_features . items ( ) :
if isinstance ( v , ( datasets . Image , VideoFrame ) ) :
self . _canonical_image_keys . add ( k )
# Cleaner: keep fps & consistent feature sets per robot type (unchanged)
cleaner = MultiLeRobotDatasetCleaner (
datasets = self . _datasets ,
repo_ids = repo_ids ,
sampling_weights = self . sampling_weights ,
datasets_repo_ids = self . datasets_repo_ids ,
min_fps = min_fps ,
max_fps = max_fps ,
)
self . _datasets = cleaner . cleaned_datasets
self . sampling_weights = cleaner . cleaned_weights
self . repo_ids = cleaner . cleaned_repo_ids
self . datasets_repo_ids = cleaner . cleaned_datasets_repo_ids
self . cumulative_sizes = cleaner . cumulative_sizes
# Meta (unchanged): we give it dim maxima; it will reshape/pad vectors
self . meta = MultiLeRobotDatasetMeta (
datasets = self . _datasets ,
repo_ids = self . repo_ids ,
keys_to_max_dim = {
ACTION : self . max_action_dim if self . max_action_dim is not None else - 1 ,
OBS_ENV_STATE : self . max_state_dim if self . max_state_dim is not None else - 1 ,
OBS_STATE : self . max_state_dim if self . max_state_dim is not None else - 1 ,
OBS_IMAGE : self . max_image_dim if self . max_image_dim is not None else - 1 ,
OBS_IMAGE_2 : self . max_image_dim if self . max_image_dim is not None else - 1 ,
OBS_IMAGE_3 : self . max_image_dim if self . max_image_dim is not None else - 1 ,
} ,
train_on_all_features = train_on_all_features ,
)
# --- NEW: track dropped (source) keys so collate won’
# Anything that we *rename away* should be considered disabled,
# otherwise downstream may expect them to exist.
self . _dropped_keys = set ( )
for rid , mapping in self . feature_keys_mapping . items ( ) :
self . _dropped_keys | = set ( mapping . keys ( ) )
# Merge with meta’
self . disabled_features = set ( self . meta . disabled_features ) | self . _dropped_keys
self . stats = self . meta . stats
# --- NEW: cache an example image shape per canonical key (lazy, filled on first use)
self . _cached_img_shape : dict [ str , torch . Size ] = { }
# ---------------------- NEW small helpers ----------------------
def _is_image_key_like ( self , key : str ) - > bool :
# A loose heuristic: rely on name OR on features later
return ( " image " in key ) or ( " cam_ " in key ) or ( " images. " in key )
def _should_ignore ( self , key : str ) - > bool :
# exact or glob-style match
for pat in self . _ignore_patterns :
if key == pat or fnmatch . fnmatch ( key , pat ) :
return True
return False
def _apply_feature_mapping ( self , item : dict , repo_id : str ) - > dict :
"""
Rename features according to feature_keys_mapping[repo_id].
- Moves tensor/image under target key.
- Drops source key if moved.
- Adds *_is_pad=False for image targets we fill/keep.
"""
mapping = self . feature_keys_mapping . get ( repo_id , { } ) or { }
if not mapping :
return item
for src , tgt in mapping . items ( ) :
if src in item :
# Move value
item [ tgt ] = item [ src ]
# Drop the source to avoid duplication
del item [ src ]
return item
def _ensure_union_image_keys ( self , item : dict ) - > dict :
"""
Ensure that every canonical image key exists.
When missing, create a zero tensor matching (B,C,H,W) or (C,H,W) of an available image.
Also add boolean mask at f " {key} _is_pad " .
"""
if not self . train_on_all_features or not self . _canonical_image_keys :
return item
# find any existing image tensor in item to copy shape/dtype
exemplar = None
for k in list ( item . keys ( ) ) :
v = item [ k ]
if torch . is_tensor ( v ) and v . ndim in ( 3 , 4 , 5 ) : # (C,H,W) or (B,C,H,W) or (B,T,C,H,W)
exemplar = v
break
# fallback to a safe 3x224x224 if nothing found
def _fallback_image ( ) :
return torch . zeros ( 3 , 224 , 224 , dtype = torch . uint8 )
for key in self . _canonical_image_keys :
if key not in item :
img = torch . zeros_like ( exemplar ) if exemplar is not None else _fallback_image ( )
item [ key ] = img
item [ f " { key } _is_pad " ] = torch . tensor ( True , dtype = torch . bool )
else :
# Add a mask saying it’
if f " { key } _is_pad " not in item :
item [ f " { key } _is_pad " ] = torch . tensor ( False , dtype = torch . bool )
return item
# ---------------------- existing API below (mostly unchanged) ----------------------
@property
def repo_id_to_index ( self ) :
""" Return a mapping from dataset repo_id to a dataset index automatically created by this class.
This index is incorporated as a data key in the dictionary returned by `__getitem__`.
"""
return { repo_id : i for i , repo_id in enumerate ( self . repo_ids ) }
@property
def repo_index_to_id ( self ) :
""" Return the inverse mapping if repo_id_to_index. """
return { v : k for k , v in self . repo_id_to_index }
@property
def fps ( self ) - > int :
""" Frames per second used during data collection.
NOTE: Fow now, this relies on a check in __init__ to make sure all sub-datasets have the same info.
"""
return self . _datasets [ 0 ] . meta . info [ " fps " ]
@property
def video ( self ) - > bool :
""" Returns True if this dataset loads video frames from mp4 files.
Returns False if it only loads images from png files.
NOTE: Fow now, this relies on a check in __init__ to make sure all sub-datasets have the same info.
"""
return self . _datasets [ 0 ] . meta . info . get ( " video " , False )
@property
def features ( self ) - > datasets . Features :
features = { }
"""
Extend native HF features with any *target* keys introduced by mapping.
We copy the source spec for targets that didn’
"""
features : dict [ str , datasets . features . Feature ] = { }
for dataset in self . _datasets :
features . update ( { k : v for k , v in dataset . hf_features . items ( ) if k not in self . disabled_features } )
for k , v in dataset . hf_features . items ( ) :
if k not in self . disabled_features :
features [ k ] = v
# Add mapped target image specs if not present yet
for rid , mapping in self . feature_keys_mapping . items ( ) :
ds = None
# find the dataset object to read feature spec for source
for _ds , _rid in zip ( self . _datasets , self . repo_ids , strict = False ) :
if _rid == rid :
ds = _ds
break
if ds is None :
continue
for src , tgt in mapping . items ( ) :
if tgt not in features and src in ds . hf_features :
features [ tgt ] = ds . hf_features [ src ]
return features
@property
def camera_keys ( self ) - > list [ str ] :
""" Keys to access image and video stream from cameras. """
keys = [ ]
for key , feats in self . features . items ( ) :
if isinstance ( feats , ( datasets . Image , VideoFrame ) ) :
@@ -1437,12 +1723,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
@property
def video_frame_keys ( self ) - > list [ str ] :
""" Keys to access video frames that requires to be decoded into images.
Note: It is empty if the dataset contains images only,
or equal to `self.cameras` if the dataset contains videos only,
or can even be a subset of `self.cameras` in a case of a mixed image/video dataset.
"""
video_frame_keys = [ ]
for key , feats in self . features . items ( ) :
if isinstance ( feats , VideoFrame ) :
@@ -1451,21 +1731,14 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
@property
def num_frames ( self ) - > int :
""" Number of samples/frames. """
return sum ( d . num_frames for d in self . _datasets )
@property
def num_episodes ( self ) - > int :
""" Number of episodes. """
return sum ( d . num_episodes for d in self . _datasets )
@property
def tolerance_s ( self ) - > float :
""" Tolerance in seconds used to discard loaded frames when their timestamps
are not close enough from the requested frames. It is only used when `delta_timestamps`
is provided or when loading video frames from mp4 files.
"""
# 1e-4 to account for possible numerical error
return 1 / self . fps - 1e-4
def __len__ ( self ) :
@@ -1474,22 +1747,83 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
def __getitem__ ( self , idx : int ) - > dict [ str , torch . Tensor ] :
if idx > = len ( self ) :
raise IndexError ( f " Index { idx } out of bounds. " )
# Determine which dataset to get an item from based on the index.
start_idx = 0
dataset_idx = 0
for dataset in self . _datasets :
if idx > = start_idx + dataset . num_frames :
start_idx + = dataset . num_frames
dataset_idx + = 1
continue
break
else :
raise AssertionError ( " We expect the loop to break out as long as the index is within bounds. " )
item = self . _datasets [ dataset_idx ] [ idx - start_idx ]
dataset_idx = np . searchsorted ( self . cumulative_sizes , idx , side = " right " ) . item ( ) - 1
local_idx = ( idx - self . cumulative_sizes [ dataset_idx ] ) . item ( )
item = self . _datasets [ dataset_idx ] [ local_idx ]
# Identify which repo this sample came from
repo_id = self . datasets_repo_ids [ dataset_idx ]
# --- NEW: apply mapping and ensure union of image keys ---
item = self . _apply_feature_mapping ( item , repo_id )
item = self . _ensure_union_image_keys ( item )
# annotate dataset index for downstream
item [ " dataset_index " ] = torch . tensor ( dataset_idx )
# Pad vector features to max dims using meta (unchanged)
item = create_padded_features ( item , self . meta . features )
# Drop any disabled (including original source keys we remapped away)
for data_key in self . disabled_features :
if data_key in item :
del item [ data_key ]
for k in IGNORED_KEYS :
if k in item :
item . pop ( k )
# Convert any datasets.Image still present to tensor
if self . image_transforms is not None :
for cam in [ k for k in item . keys ( ) if self . _is_image_key_like ( k ) ] :
val = item [ cam ]
if not torch . is_tensor ( val ) :
item [ cam ] = self . image_transforms ( val )
# 🔑 Pad actions if too short
if " actions " in item and self . max_action_dim is not None :
act = item [ " actions " ]
if act . shape [ - 1 ] < self . max_action_dim :
pad_len = self . max_action_dim - act . shape [ - 1 ]
item [ " actions " ] = torch . cat ( [ act , torch . zeros ( pad_len , dtype = act . dtype ) ] , dim = - 1 )
item [ " actions_padding_mask " ] = torch . cat (
[ torch . zeros_like ( act , dtype = torch . bool ) , torch . ones ( pad_len , dtype = torch . bool ) ] ,
dim = - 1 ,
)
# pad obs_state if too short
if " obs_state " in item and self . max_state_dim is not None :
st = item [ " obs_state " ]
if st . shape [ - 1 ] < self . max_state_dim :
pad_len = self . max_state_dim - st . shape [ - 1 ]
item [ " obs_state " ] = torch . cat ( [ st , torch . zeros ( pad_len , dtype = st . dtype ) ] , dim = - 1 )
item [ " obs_state_padding_mask " ] = torch . cat (
[ torch . zeros_like ( st , dtype = torch . bool ) , torch . ones ( pad_len , dtype = torch . bool ) ] ,
dim = - 1 ,
)
# actions
if " actions " in item and self . max_action_dim is not None :
act = item [ " actions " ]
if act . shape [ - 1 ] < self . max_action_dim :
pad_len = self . max_action_dim - act . shape [ - 1 ]
item [ " actions " ] = torch . cat ( [ act , torch . zeros ( pad_len , dtype = act . dtype ) ] , dim = - 1 )
mask = torch . cat (
[ torch . zeros_like ( act , dtype = torch . bool ) , torch . ones ( pad_len , dtype = torch . bool ) ] ,
dim = - 1 ,
)
else :
mask = torch . zeros ( self . max_action_dim , dtype = torch . bool ) # 👈 all False if no padding
item [ " actions_padding_mask " ] = mask
# obs state
if " obs_state " in item and self . max_state_dim is not None :
st = item [ " obs_state " ]
if st . shape [ - 1 ] < self . max_state_dim :
pad_len = self . max_state_dim - st . shape [ - 1 ]
item [ " obs_state " ] = torch . cat ( [ st , torch . zeros ( pad_len , dtype = st . dtype ) ] , dim = - 1 )
mask = torch . cat (
[ torch . zeros_like ( st , dtype = torch . bool ) , torch . ones ( pad_len , dtype = torch . bool ) ] ,
dim = - 1 ,
)
else :
mask = torch . zeros ( self . max_state_dim , dtype = torch . bool ) # 👈 always add mask
item [ " obs_state_padding_mask " ] = mask
return item
@@ -1506,3 +1840,149 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
f " Transformations: { self . image_transforms } , \n "
f " ) "
)
def keep_datasets_with_the_same_features_per_robot_type ( ls_datasets : list ) - > list :
"""
Filters datasets to only keep those with consistent feature shapes per robot type.
Args:
ls_datasets (List): List of datasets, each with a `meta.info[ ' robot_type ' ]`
and `meta.episodes_stats` dictionary.
Returns:
List: Filtered list of datasets with consistent feature shapes.
"""
robot_types = { ds . meta . info [ " robot_type " ] for ds in ls_datasets }
datasets_to_remove = set ( )
for robot_type in robot_types :
# Collect all stats dicts for this robot type
stats_list = [
ep_stats
for ds in ls_datasets
if ds . meta . info [ " robot_type " ] == robot_type
for ep_stats in episode_stats_values ( ds . meta )
]
if not stats_list :
continue
# Determine the most common shape for each key
all_keys = { key for stats in stats_list for key in stats }
for ds in ls_datasets :
if ds . meta . info [ " robot_type " ] != robot_type :
continue
for key in all_keys :
shape_counter = defaultdict ( int )
for stats in stats_list :
value = stats . get ( key )
if (
value and " mean " in value and isinstance ( value [ " mean " ] , ( torch . Tensor , np . ndarray ) )
) : # FIXME(mshukor): check all stats; min, mean, max
shape_counter [ value [ " mean " ] . shape ] + = 1
if not shape_counter :
continue
# Identify the most frequent shape
main_shape = max ( shape_counter , key = shape_counter . get )
# Flag datasets that don't match the main shape
# for ds in ls_datasets:
first_ep_stats = next ( iter ( episode_stats_values ( ds . meta ) ) , None )
if not first_ep_stats :
continue
value = first_ep_stats . get ( key )
if (
value
and " mean " in value
and isinstance ( value [ " mean " ] , ( torch . Tensor , np . ndarray ) )
and value [ " mean " ] . shape != main_shape
) :
datasets_to_remove . add ( ds )
break
# Filter out inconsistent datasets
datasets_maks = [ ds not in datasets_to_remove for ds in ls_datasets ]
filtered_datasets = [ ds for ds in ls_datasets if ds not in datasets_to_remove ]
print (
f " Keeping { len ( filtered_datasets ) } datasets. Removed { len ( datasets_to_remove ) } inconsistent ones. Inconsistent datasets: \n { datasets_to_remove } "
)
return filtered_datasets , datasets_maks
def aggregate_stats_per_robot_type ( ls_datasets ) - > dict [ str , dict [ str , torch . Tensor ] ] :
""" Aggregate stats of multiple LeRobot datasets into multiple set of stats per robot type.
The final stats will have the union of all data keys from each of the datasets.
The final stats will have the union of all data keys from each of the datasets. For instance:
- new_max = max(max_dataset_0, max_dataset_1, ...)
- new_min = min(min_dataset_0, min_dataset_1, ...)
- new_mean = (mean of all data)
- new_std = (std of all data)
"""
robot_types = { ds . meta . info [ " robot_type " ] for ds in ls_datasets }
stats = { robot_type : { } for robot_type in robot_types }
for robot_type in robot_types :
robot_type_datasets = [ ]
for ds in ls_datasets :
if ds . meta . info [ " robot_type " ] == robot_type :
robot_type_datasets . extend ( list ( episode_stats_values ( ds . meta ) ) )
# robot_type_datasets = [list(ds.episodes_stats.values()) for ds in ls_datasets if ds.meta.info["robot_type"] == robot_type]
stat = aggregate_stats ( robot_type_datasets )
stats [ robot_type ] = stat
return stats
def reshape_features_to_max_dim ( features : dict , reshape_dim : int = - 1 , keys_to_max_dim : dict = { } ) - > dict :
""" Reshape features to have a maximum dimension of `max_dim`. """
reshaped_features = { }
for key in features :
if key in keys_to_max_dim and keys_to_max_dim [ key ] is not None :
reshaped_features [ key ] = features [ key ]
shape = list ( features [ key ] [ " shape " ] )
if any ( [ k in key for k in [ OBS_IMAGE , OBS_IMAGE_2 , OBS_IMAGE_3 ] ] ) : # Assume square images
shape [ - 3 ] = keys_to_max_dim [ key ]
shape [ - 2 ] = keys_to_max_dim [ key ]
else :
shape [ reshape_dim ] = keys_to_max_dim [ key ]
reshaped_features [ key ] [ " shape " ] = tuple ( shape )
else :
reshaped_features [ key ] = features [ key ]
return reshaped_features
def create_padded_features ( item : dict , features : dict = { } ) :
for key , ft in features . items ( ) :
if any ( [ k in key for k in [ " cam " , " effort " , " absolute " ] ] ) : # FIXME(mshukor): temporary hack
continue
shape = ft [ " shape " ]
if len ( shape ) == 3 : # images to torch format (C, H, W)
shape = ( shape [ 2 ] , shape [ 0 ] , shape [ 1 ] )
if len ( shape ) == 1 and shape [ 0 ] == 1 : # ft with shape are actually tensor(ele)
shape = [ ]
if key not in item :
dtype = str_to_torch_dtype ( ft [ " dtype " ] )
item [ key ] = torch . zeros ( shape , dtype = dtype )
item [ f " { key } _padding_mask " ] = torch . tensor ( 0 , dtype = torch . int64 )
if " image " in key : # FIXME(mshukor): support other observations
item [ f " { key } _is_pad " ] = torch . BoolTensor ( [ False ] )
else :
item [ f " { key } _padding_mask " ] = torch . tensor ( 1 , dtype = torch . int64 )
return item
def str_to_torch_dtype ( dtype_str ) :
""" Convert a dtype string to a torch dtype. """
mapping = {
" float32 " : torch . float32 ,
" int64 " : torch . int64 ,
" int16 " : torch . int16 ,
" bool " : torch . bool ,
" video " : torch . float32 , # Assuming video is stored as uint8 images
}
return mapping . get ( dtype_str , torch . float32 ) # Default to float32
def episode_stats_values ( meta ) :
episodes = meta . episodes . to_pandas ( ) . to_dict ( orient = " records " )
return [
{ k : v for k , v in ep . items ( ) if isinstance ( v , dict ) and " mean " in v }
for ep in episodes
]
Jade Choghari