diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 2fb68dca1..d1e612b60 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -234,6 +234,8 @@ def merge_datasets( datasets: list[LeRobotDataset], output_repo_id: str, output_dir: str | Path | None = None, + data_files_size_in_mb: float | None = None, + video_files_size_in_mb: float | None = None, ) -> LeRobotDataset: """Merge multiple LeRobotDatasets into a single dataset. @@ -257,6 +259,8 @@ def merge_datasets( aggr_repo_id=output_repo_id, roots=roots, aggr_root=output_dir, + data_files_size_in_mb=data_files_size_in_mb, + video_files_size_in_mb=video_files_size_in_mb, ) merged_dataset = LeRobotDataset( @@ -747,11 +751,11 @@ def _copy_and_reindex_videos( f"videos/{video_key}/to_timestamp" ] else: - # Build list of time ranges to keep, in sorted order. - sorted_keep_episodes = sorted(episodes_in_file, key=lambda x: episode_mapping[x]) + # Episodes are already in order by old episode index (from episode_mapping iteration), + # which equals from_timestamp order since episodes are created sequentially. episodes_to_keep_ranges: list[tuple[float, float]] = [] - for old_idx in sorted_keep_episodes: + for old_idx in episodes_in_file: src_ep = src_dataset.meta.episodes[old_idx] from_ts = src_ep[f"videos/{video_key}/from_timestamp"] to_ts = src_ep[f"videos/{video_key}/to_timestamp"] @@ -781,7 +785,7 @@ def _copy_and_reindex_videos( ) cumulative_ts = 0.0 - for old_idx in sorted_keep_episodes: + for old_idx in episodes_in_file: new_idx = episode_mapping[old_idx] src_ep = src_dataset.meta.episodes[old_idx] ep_length = src_ep["length"] diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index d9d4b22d0..71e715f49 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -557,7 +557,7 @@ class LeRobotDataset(torch.utils.data.Dataset): episodes: list[int] | None = None, image_transforms: Callable | None = None, delta_timestamps: dict[str, list[float]] | None = None, - tolerance_s: float = 1e-4, + tolerance_s: float = 1e-2, revision: str | None = None, force_cache_sync: bool = False, download_videos: bool = True, diff --git a/src/lerobot/optim/optimizers.py b/src/lerobot/optim/optimizers.py index 5120f828c..dde118f39 100644 --- a/src/lerobot/optim/optimizers.py +++ b/src/lerobot/optim/optimizers.py @@ -81,10 +81,14 @@ class AdamWConfig(OptimizerConfig): eps: float = 1e-8 weight_decay: float = 1e-2 grad_clip_norm: float = 10.0 + fused: bool = False def build(self, params: dict) -> torch.optim.Optimizer: kwargs = asdict(self) kwargs.pop("grad_clip_norm") + # Fused optimizer only works on CUDA + if kwargs.get("fused") and not torch.cuda.is_available(): + kwargs["fused"] = False return torch.optim.AdamW(params, **kwargs) diff --git a/src/lerobot/policies/act/configuration_act.py b/src/lerobot/policies/act/configuration_act.py index 6f6c1c4be..f9e6fd35a 100644 --- a/src/lerobot/policies/act/configuration_act.py +++ b/src/lerobot/policies/act/configuration_act.py @@ -136,6 +136,7 @@ class ACTConfig(PreTrainedConfig): optimizer_lr: float = 1e-5 optimizer_weight_decay: float = 1e-4 optimizer_lr_backbone: float = 1e-5 + optimizer_fused: bool = False # Use CUDA fused AdamW kernel def __post_init__(self): super().__post_init__() @@ -164,6 +165,7 @@ class ACTConfig(PreTrainedConfig): return AdamWConfig( lr=self.optimizer_lr, weight_decay=self.optimizer_weight_decay, + fused=self.optimizer_fused, ) def get_scheduler_preset(self) -> None: diff --git a/src/lerobot/policies/groot/configuration_groot.py b/src/lerobot/policies/groot/configuration_groot.py index 8002c69ea..167d95af5 100644 --- a/src/lerobot/policies/groot/configuration_groot.py +++ b/src/lerobot/policies/groot/configuration_groot.py @@ -94,6 +94,7 @@ class GrootConfig(PreTrainedConfig): optimizer_betas: tuple[float, float] = (0.95, 0.999) optimizer_eps: float = 1e-8 optimizer_weight_decay: float = 1e-5 + optimizer_fused: bool = False # Use CUDA fused AdamW kernel warmup_ratio: float = 0.05 use_bf16: bool = True @@ -174,6 +175,7 @@ class GrootConfig(PreTrainedConfig): betas=self.optimizer_betas, eps=self.optimizer_eps, weight_decay=self.optimizer_weight_decay, + fused=self.optimizer_fused, ) def get_scheduler_preset(self) -> CosineDecayWithWarmupSchedulerConfig: diff --git a/src/lerobot/policies/pi0/configuration_pi0.py b/src/lerobot/policies/pi0/configuration_pi0.py index 33753e0b2..c06cd98b0 100644 --- a/src/lerobot/policies/pi0/configuration_pi0.py +++ b/src/lerobot/policies/pi0/configuration_pi0.py @@ -74,6 +74,7 @@ class PI0Config(PreTrainedConfig): gradient_checkpointing: bool = False # Enable gradient checkpointing for memory optimization compile_model: bool = False # Whether to use torch.compile for model optimization compile_mode: str = "max-autotune" # Torch compile mode + optimizer_fused: bool = False # Use CUDA fused AdamW kernel device: str | None = None # Device to use for the model (None = auto-detect) # Optimizer settings: see openpi `AdamW`` @@ -141,6 +142,7 @@ class PI0Config(PreTrainedConfig): eps=self.optimizer_eps, weight_decay=self.optimizer_weight_decay, grad_clip_norm=self.optimizer_grad_clip_norm, + fused=self.optimizer_fused, ) def get_scheduler_preset(self): diff --git a/src/lerobot/policies/pi05/configuration_pi05.py b/src/lerobot/policies/pi05/configuration_pi05.py index 7bdce70dd..6e685697e 100644 --- a/src/lerobot/policies/pi05/configuration_pi05.py +++ b/src/lerobot/policies/pi05/configuration_pi05.py @@ -74,6 +74,7 @@ class PI05Config(PreTrainedConfig): gradient_checkpointing: bool = False # Enable gradient checkpointing for memory optimization compile_model: bool = False # Whether to use torch.compile for model optimization compile_mode: str = "max-autotune" # Torch compile mode + optimizer_fused: bool = False # Use CUDA fused AdamW kernel device: str | None = None # Device to use for the model (None = auto-detect) # Optimizer settings: see openpi `AdamW` @@ -141,6 +142,7 @@ class PI05Config(PreTrainedConfig): eps=self.optimizer_eps, weight_decay=self.optimizer_weight_decay, grad_clip_norm=self.optimizer_grad_clip_norm, + fused=self.optimizer_fused, ) def get_scheduler_preset(self): diff --git a/src/lerobot/policies/smolvla/configuration_smolvla.py b/src/lerobot/policies/smolvla/configuration_smolvla.py index c32c8a60e..eea7299a1 100644 --- a/src/lerobot/policies/smolvla/configuration_smolvla.py +++ b/src/lerobot/policies/smolvla/configuration_smolvla.py @@ -79,6 +79,7 @@ class SmolVLAConfig(PreTrainedConfig): optimizer_eps: float = 1e-8 optimizer_weight_decay: float = 1e-10 optimizer_grad_clip_norm: float = 10 + optimizer_fused: bool = False scheduler_warmup_steps: int = 1_000 scheduler_decay_steps: int = 30_000 @@ -136,6 +137,7 @@ class SmolVLAConfig(PreTrainedConfig): eps=self.optimizer_eps, weight_decay=self.optimizer_weight_decay, grad_clip_norm=self.optimizer_grad_clip_norm, + fused=self.optimizer_fused, ) def get_scheduler_preset(self):