| """Native LaWAM samples from immutable RoboCerebra HDF5, with resumable anchors.""" |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import sys |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import DataLoader, Dataset, Sampler |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT)) |
| sys.path.insert(0, str(ROOT / "LaWAM-original")) |
| from training.robocerebra_train100.dataset import RoboCerebraHDF5Dataset |
| from starVLA.dataloader.gr00t_lerobot.transform.state_action import Normalizer |
|
|
| MASK64 = (1 << 64) - 1 |
|
|
|
|
| def frame_for_ordinal(ordinal: int, frame_count: int, seed: int) -> int: |
| """Stateless SplitMix64 frame draw; independent of rank, workers and prefetch.""" |
| value = (int(ordinal) + int(seed) + 0x9E3779B97F4A7C15) & MASK64 |
| value = ((value ^ (value >> 30)) * 0xBF58476D1CE4E5B9) & MASK64 |
| value = ((value ^ (value >> 27)) * 0x94D049BB133111EB) & MASK64 |
| value ^= value >> 31 |
| return value % int(frame_count) |
|
|
|
|
| class GlobalOrdinalBatchSampler(Sampler): |
| """Partition the same ordered 512 anchors across any compatible DDP topology.""" |
| def __init__(self, *, start_step, stop_step, global_batch, microbatch, rank, world_size): |
| self.start_step, self.stop_step = int(start_step), int(stop_step) |
| self.global_batch, self.microbatch = int(global_batch), int(microbatch) |
| self.rank, self.world_size = int(rank), int(world_size) |
| if self.global_batch % (self.microbatch * self.world_size): |
| raise ValueError("Global batch must divide microbatch * world_size exactly") |
| if not 0 <= self.rank < self.world_size or self.start_step > self.stop_step: |
| raise ValueError("Invalid rank or optimizer-step bounds") |
| self.accumulation = self.global_batch // (self.microbatch * self.world_size) |
|
|
| def __iter__(self): |
| for step in range(self.start_step, self.stop_step): |
| for microstep in range(self.accumulation): |
| first = step * self.global_batch + (microstep * self.world_size + self.rank) * self.microbatch |
| yield list(range(first, first + self.microbatch)) |
|
|
| def __len__(self): |
| return (self.stop_step - self.start_step) * self.accumulation |
|
|
|
|
| class NativeActions: |
| def __init__(self, stats_path): |
| self.path = Path(stats_path) |
| payload = self.path.read_bytes() |
| self.sha256 = hashlib.sha256(payload).hexdigest() |
| self.stats = json.loads(payload) |
| raw = self.stats["raw_action"] |
| self.minimum = np.asarray(raw["min"], dtype=np.float32) |
| self.maximum = np.asarray(raw["max"], dtype=np.float32) |
| if self.minimum.shape != (7,) or self.maximum.shape != (7,): |
| raise ValueError("Expected current-dataset seven-dimensional action statistics") |
| if not np.isfinite(self.minimum).all() or not np.isfinite(self.maximum).all(): |
| raise ValueError("Action normalization statistics must be finite") |
| if np.any(self.maximum < self.minimum): |
| raise ValueError("Action normalization bounds are reversed") |
| self.pose = Normalizer("min_max", {"min": self.minimum[:6].tolist(), "max": self.maximum[:6].tolist()}) |
| self.gripper = Normalizer("binary", {}, binary_threshold=0.5) |
|
|
| def forward(self, raw): |
| raw = torch.as_tensor(np.asarray(raw, dtype=np.float32).copy()) |
| if raw.ndim != 2 or raw.shape[1] != 7 or not torch.isfinite(raw).all(): |
| raise ValueError("Expected finite [H,7] raw action") |
| if not torch.all((raw[:, 6] == -1) | (raw[:, 6] == 1)): |
| raise ValueError("RoboCerebra raw gripper must retain signed -1/+1 commands") |
| return torch.cat((self.pose.forward(raw[:, :6]), self.gripper.forward(raw[:, 6:])), dim=1) |
|
|
| def inverse(self, encoded, *, clip=False): |
| """Decode [...,7] predictions to RoboCerebra controller units and signs.""" |
| encoded = torch.as_tensor(encoded, dtype=torch.float32) |
| if encoded.shape[-1] != 7: |
| raise ValueError("Expected seven action coordinates for controller decoding") |
| |
| value = encoded.detach().cpu().clone() |
| if clip: |
| value[..., :6].clamp_(-1, 1) |
| pose = self.pose.inverse(value[..., :6]) |
| signed = 2 * (value[..., 6:] > 0.5).to(value.dtype) - 1 |
| return torch.cat((pose, signed), dim=-1).to(encoded.device) |
|
|
| def export_statistics(self, frames, episodes): |
| action = dict(self.stats["lawam_binary_gripper_action"]) |
| action["mask"] = [True] * 6 + [False] |
| return {"franka": {"action": action, "num_transitions": int(frames), "num_trajectories": int(episodes)}} |
|
|
|
|
| class LaWAMDataset(Dataset): |
| def __init__(self, *, data_root, index_path, stats_path, seed=2026, sampled=True, max_open_files=8): |
| self.raw = RoboCerebraHDF5Dataset( |
| data_root=data_root, index_path=index_path, horizon=8, image_size=None, |
| include_future_image=True, future_offset=7, max_open_files=max_open_files, |
| ) |
| self.normalization = NativeActions(stats_path) |
| self.seed, self.sampled = int(seed), bool(sampled) |
| self.index = json.loads(Path(index_path).read_text()) |
| self.case_ids = {entry["case_id"] for entry in self.index["episodes"]} |
|
|
| def __len__(self): |
| return len(self.raw) |
|
|
| def frame_index(self, ordinal): |
| return frame_for_ordinal(ordinal, len(self.raw), self.seed) if self.sampled else int(ordinal) |
|
|
| def __getitem__(self, ordinal): |
| frame = self.frame_index(ordinal) |
| raw = self.raw[frame] |
| current = np.asarray(raw["images"]["agentview"]) |
| future = np.asarray(raw["images"]["future_agentview"]) |
| wrist = np.asarray(raw["images"]["eye_in_hand"]) |
| for image in (current, future, wrist): |
| if image.dtype != np.uint8 or image.shape != (256, 256, 3): |
| raise ValueError(f"Expected unchanged rendered256 RGB, got {image.shape}/{image.dtype}") |
| action = self.normalization.forward(raw["actions"]) |
| if action.shape != (8, 7): |
| raise ValueError(f"Expected native8-action chunk, got {action.shape}") |
| |
| |
| return { |
| "primary_videos": torch.from_numpy(np.stack((current, future)).transpose(0, 3, 1, 2).copy()).unsqueeze(0), |
| "wrist_images": torch.from_numpy(wrist.transpose(2, 0, 1).copy()).unsqueeze(0), |
| "state": torch.zeros((1, 7), dtype=torch.float32), |
| "action": action, "lang": str(raw["instruction"]), |
| "embodiment_id": 25, "action_hz": 20.0, |
| "audit": {"global_ordinal": int(ordinal), "frame_anchor": frame, |
| "source": raw["metadata"], "genuine_action_mask": np.asarray(raw["action_valid_mask"]).tolist()}, |
| } |
|
|
|
|
| def worker_init(_worker_id): |
| for name in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS"): |
| os.environ[name] = "1" |
| torch.set_num_threads(1) |
|
|
|
|
| def build_dataloaders(cfg, *, start_step, world_size, rank): |
| from starVLA.dataloader import _build_latent_world_collator |
| from starVLA.model.framework.latent_world.config_builder import LatentWorldPolicyConfigBuilder |
| data = cfg.datasets.vla_data |
| policy = LatentWorldPolicyConfigBuilder(cfg).build() |
| if bool(policy.flow_cfg.use_state): |
| raise ValueError("This baseline preserves native LaWAM use_state:false") |
| train = LaWAMDataset(data_root=data.data_root_dir, index_path=data.index_path, |
| stats_path=data.action_stats_path, seed=cfg.seed, sampled=True) |
| val = LaWAMDataset(data_root=data.val_data_root_dir, index_path=data.val_index_path, |
| stats_path=data.action_stats_path, seed=cfg.seed, sampled=False) |
| if len(train.case_ids) != 100 or len(val.case_ids) != 20 or train.case_ids & val.case_ids: |
| raise ValueError("Require disjoint100 training and20 validation parent trajectories") |
| sampler = GlobalOrdinalBatchSampler(start_step=start_step, stop_step=cfg.trainer.max_train_steps, |
| global_batch=cfg.trainer.global_batch_size, |
| microbatch=data.per_device_batch_size, rank=rank, world_size=world_size) |
| workers = int(data.num_workers) |
| common = dict(num_workers=workers, pin_memory=True) |
| if workers: |
| common.update(prefetch_factor=int(data.prefetch_factor), persistent_workers=True, |
| worker_init_fn=worker_init, in_order=True) |
| train_loader = DataLoader(train, batch_sampler=sampler, |
| collate_fn=_build_latent_world_collator(cfg, policy_cfg=policy, training=True), **common) |
| |
| total = int(cfg.trainer.eval_global_anchor_count) |
| if total % (int(data.per_device_batch_size) * int(world_size)): |
| raise ValueError("Fixed validation anchor count must form complete equal-rank batches") |
| indices = np.linspace(0, len(val)-1, num=min(total, len(val)), dtype=np.int64).tolist()[rank::world_size] |
| val_loader = DataLoader(val, batch_size=int(data.per_device_batch_size), sampler=indices, |
| collate_fn=_build_latent_world_collator(cfg, policy_cfg=policy, training=False), **common) |
| if rank == 0: |
| out = Path(cfg.output_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| statistics = train.normalization.export_statistics(len(train), train.index["episode_count"]) |
| (out / "dataset_statistics.json").write_text(json.dumps(statistics, indent=2) + "\n") |
| return train_loader, val_loader |
|
|
|
|
| class DeviceLoader: |
| """Already sharded sampler: move batches without Accelerate sharding a second time.""" |
| def __init__(self, loader, device): |
| self.loader, self.device = loader, device |
| self.dataset = loader.dataset |
|
|
| def __len__(self): |
| return len(self.loader) |
|
|
| def __iter__(self): |
| from accelerate.utils import send_to_device |
| for batch in self.loader: |
| yield send_to_device(batch, self.device, non_blocking=True) |
|
|
| def close(self): |
| iterator = getattr(self.loader, "_iterator", None) |
| if iterator is not None: |
| iterator._shutdown_workers() |
|
|
