"""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") # Upstream Normalizer owns CPU statistics, as in its native dataset path. 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}") # Native LIBERO repeats terminal actions and counts all8 labels as valid. # The source's genuine-frame mask remains audit metadata, not a new loss rule. 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) # Evenly spread a fixed validation anchor set, with no overlap across ranks. 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()