| """Image trajectory dataset for world-model training.""" |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
|
|
| from driftwm.data.generate import ID_TO_BOAT |
| from experiments.shared.src.vision.clean_renderer import render_clean_boat_array |
|
|
|
|
| class ImageTrajectoryDataset(Dataset): |
| def __init__( |
| self, |
| source_npz: str, |
| history_len: int, |
| horizon: int, |
| episodes: int, |
| max_windows: int, |
| seed: int, |
| episode_start: int = 0, |
| image_size: int = 160, |
| visual_scale: float = 2.5, |
| render_images: bool = True, |
| return_origin: bool = False, |
| return_aux: bool = False, |
| ): |
| src = np.load(source_npz, allow_pickle=False) |
| start = int(episode_start) |
| end = start + int(episodes) |
| if end > int(src["obs"].shape[0]): |
| raise ValueError(f"requested episodes [{start}, {end}) but dataset only has {src['obs'].shape[0]}") |
| self.obs = src["obs"][start:end].astype(np.float32) |
| self.actions = src["actions"][start:end].astype(np.float32) |
| self.flow_type_ids = src["flow_type_ids"][start:end].astype(np.int64) |
| self.traj_type_ids = src["traj_type_ids"][start:end].astype(np.int64) |
| self.boat_ids = src["boat_ids"][start:end].astype(np.int64) |
| self.states = src["states"][start:end].astype(np.float32) |
| self.image_size = int(image_size) |
| self.visual_scale = float(visual_scale) |
| self.render_images = bool(render_images) |
| self.history_len = int(history_len) |
| self.horizon = int(horizon) |
| steps = self.actions.shape[1] |
| indices = [(ep, t) for ep in range(episodes) for t in range(self.history_len - 1, steps - self.horizon)] |
| rng = np.random.default_rng(seed) |
| selected = rng.choice(len(indices), size=min(max_windows, len(indices)), replace=False) |
| self.indices = [indices[int(i)] for i in selected] |
| self.return_origin = bool(return_origin) |
| self.return_aux = bool(return_aux) |
|
|
| def __len__(self) -> int: |
| return len(self.indices) |
|
|
| def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| ep, t = self.indices[idx] |
| state_hist = self.states[ep, t - self.history_len + 1 : t + 1, :6].copy() |
| if self.render_images: |
| boat = ID_TO_BOAT[int(self.boat_ids[ep])] |
| rendered = [ |
| render_clean_boat_array(state, boat, image_size=self.image_size, visual_scale=self.visual_scale) |
| for state in state_hist |
| ] |
| observation_hist = np.transpose(np.stack(rendered, axis=0), (0, 3, 1, 2)) |
| else: |
| observation_hist = state_hist |
| padded_actions = np.zeros((self.actions.shape[1] + 1, self.actions.shape[2]), dtype=np.float32) |
| padded_actions[1:] = self.actions[ep] |
| action_hist = padded_actions[t - self.history_len + 1 : t + 1].copy() |
| future_actions = self.actions[ep, t : t + self.horizon].copy() |
| targets = self.obs[ep, t + 1 : t + 1 + self.horizon].copy() |
| sample = ( |
| torch.from_numpy(observation_hist), |
| torch.from_numpy(action_hist), |
| torch.from_numpy(future_actions), |
| torch.from_numpy(targets), |
| ) |
| if self.return_aux: |
| return ( |
| *sample, |
| torch.from_numpy(self.obs[ep, t].copy()), |
| torch.from_numpy(self.obs[ep, t - 1].copy()), |
| torch.tensor(self.flow_type_ids[ep], dtype=torch.long), |
| torch.tensor(self.boat_ids[ep], dtype=torch.long), |
| ) |
| if self.return_origin: |
| return (*sample, torch.from_numpy(self.obs[ep, t].copy())) |
| return sample |
|
|
