| |
| import json |
| import os |
| from accelerate.logging import get_logger |
| import numpy as np |
| from torch.utils.data import DataLoader |
| import torch.distributed as dist |
| from pathlib import Path |
| from starVLA.dataloader.infoflow_sampler import StatelessDrawSampler |
|
|
| logger = get_logger(__name__) |
|
|
| def save_dataset_statistics(dataset_statistics, run_dir): |
| """Saves a `dataset_statistics.json` file.""" |
| out_path = run_dir / "dataset_statistics.json" |
| with open(out_path, "w") as f_json: |
| for _, stats in dataset_statistics.items(): |
| for k in stats["action"].keys(): |
| if isinstance(stats["action"][k], np.ndarray): |
| stats["action"][k] = stats["action"][k].tolist() |
| if "proprio" in stats: |
| for k in stats["proprio"].keys(): |
| if isinstance(stats["proprio"][k], np.ndarray): |
| stats["proprio"][k] = stats["proprio"][k].tolist() |
| if "num_trajectories" in stats: |
| if isinstance(stats["num_trajectories"], np.ndarray): |
| stats["num_trajectories"] = stats["num_trajectories"].item() |
| if "num_transitions" in stats: |
| if isinstance(stats["num_transitions"], np.ndarray): |
| stats["num_transitions"] = stats["num_transitions"].item() |
| json.dump(dataset_statistics, f_json, indent=2) |
| logger.info(f"Saved dataset statistics file at path {out_path}") |
|
|
|
|
|
|
| def build_dataloader(cfg, dataset_py="lerobot_datasets_oxe"): |
|
|
| if dataset_py == "lerobot_datasets": |
| from starVLA.dataloader.lerobot_datasets import get_vla_dataset, collate_fn |
| vla_dataset_cfg = cfg.datasets.vla_data |
|
|
| vla_dataset = get_vla_dataset( |
| data_cfg=vla_dataset_cfg, |
| balance_dataset_weights=vla_dataset_cfg.get("balance_dataset_weights", False), |
| balance_trajectory_weights=vla_dataset_cfg.get("balance_trajectory_weights", False), |
| ) |
|
|
| traceable_sampling = bool(vla_dataset_cfg.get("traceable_sampling", False)) |
| per_device_batch_size = int(cfg.datasets.vla_data.per_device_batch_size) |
| world_size = ( |
| dist.get_world_size() |
| if dist.is_available() and dist.is_initialized() |
| else int(vla_dataset_cfg.get("expected_world_size", 1)) |
| ) |
| sampler = None |
| if traceable_sampling: |
| expected_world_size = int(vla_dataset_cfg.expected_world_size) |
| if world_size != expected_world_size: |
| raise RuntimeError( |
| f"Traceable sampling requires world size {expected_world_size}, " |
| f"but distributed runtime reports {world_size}" |
| ) |
| gradient_accumulation_steps = int( |
| cfg.trainer.gradient_accumulation_steps |
| ) |
| global_update_batch = ( |
| per_device_batch_size |
| * world_size |
| * gradient_accumulation_steps |
| ) |
| usable_draws = len(vla_dataset) - ( |
| len(vla_dataset) % global_update_batch |
| ) |
| if usable_draws <= 0: |
| raise ValueError( |
| f"Dataset length {len(vla_dataset)} is smaller than one global " |
| f"optimizer update ({global_update_batch} draws)" |
| ) |
| sampler = StatelessDrawSampler(usable_draws) |
| num_workers = int(vla_dataset_cfg.get("num_workers", 16)) |
| dataloader_kwargs = dict( |
| dataset=vla_dataset, |
| batch_size=per_device_batch_size, |
| collate_fn=collate_fn, |
| sampler=sampler, |
| num_workers=num_workers, |
| pin_memory=True, |
| persistent_workers=num_workers > 0, |
| drop_last=traceable_sampling, |
| ) |
| if num_workers > 0: |
| dataloader_kwargs["prefetch_factor"] = int(vla_dataset_cfg.get("prefetch_factor", 4)) |
| vla_train_dataloader = DataLoader(**dataloader_kwargs) |
|
|
| rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 |
| if rank == 0: |
| output_dir = Path(cfg.output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| vla_dataset.save_dataset_statistics(output_dir / "dataset_statistics.json") |
| if vla_dataset.valid_anchor_manifest is not None: |
| manifest_path = output_dir / "valid_anchor_manifest.json" |
| vla_dataset.valid_anchor_manifest.save(manifest_path) |
| return vla_train_dataloader |
| raise ValueError(f"Unsupported dataset loader: {dataset_py}") |
|
|
