File size: 5,387 Bytes
872cf4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """Offline goal-conditioned samples drawn from the cached PushT latents.
A sample is a 3-frame latent history, the two action blocks that connect
those frames, a goal latent taken from 1-``max_offset`` transitions ahead in
the *same* episode, and the real action block leaving the current frame (used
only to fit the behavior-density model, never as a controller target).
Everything is in latent space, so this never touches the 46 GB image file.
"""
import json
from pathlib import Path
import numpy as np
import torch
from torch.utils.data import Dataset
NUM_CONTEXT = 3
FRAMESKIP = 5
class LatentGoalDataset(Dataset):
"""Latent history/goal pairs sampled from cached expert trajectories.
Args:
root: Directory written by ``scripts/encode_latents.py``.
max_offset: Largest goal distance, in world-model transitions. Set by
the training curriculum (1-2, then 1-3, then 1-5).
episodes: Episode indices this split may draw from.
horizon: Plan horizon, reserved so a full plan stays inside the episode.
"""
def __init__(
self, root, max_offset=5, episodes=None, horizon=5, in_memory=True
):
root = Path(root)
stats = json.loads((root / 'stats.json').read_text())
self.latent_dim = stats['latent_dim']
# 2.34M x 192 fp16 is ~0.9 GB, so keeping it resident beats paging a
# memmap with random access from every dataloader worker.
mode = None if in_memory else 'r'
self.latents = np.load(root / 'latents.npy', mmap_mode=mode)
self.actions = np.load(root / 'actions.npy', mmap_mode=mode)
self.lengths = np.load(root / 'ep_len.npy')
self.offsets = np.load(root / 'ep_offset.npy')
self.action_mean = np.asarray(stats['action_mean'], dtype=np.float32)
self.action_std = np.asarray(stats['action_std'], dtype=np.float32)
self.horizon = horizon
self.set_max_offset(max_offset)
if episodes is None:
episodes = np.arange(len(self.lengths))
self.episodes = np.asarray(episodes)
# Flat (episode, start) index as arrays — a Python list of 2M tuples
# gets re-pickled into every worker and dominates startup cost.
# A clip needs its N context frames *and* at least one further
# transition to place a goal in; without that reserve the last clip of
# each episode silently takes its goal from the next episode.
counts = (self.lengths[self.episodes] - NUM_CONTEXT * FRAMESKIP).clip(
min=0
)
self.clip_episode = np.repeat(self.episodes, counts)
self.clip_start = np.concatenate(
[np.arange(c) for c in counts if c > 0]
).astype(np.int64)
def set_max_offset(self, max_offset):
"""Update the curriculum's furthest goal, in world-model transitions."""
self.max_offset = max_offset
def __len__(self):
return len(self.clip_start)
def _block(self, base, local_step):
"""Normalized action block leaving the frame at ``local_step``.
Matches ``swm.data.Dataset.__getitem__``: pixels are strided by
frameskip while actions are not, so the block leaving frame ``k`` is
the raw actions ``[start + 5k, start + 5k + 5)``.
"""
lo = base + local_step
raw = np.asarray(self.actions[lo : lo + FRAMESKIP], dtype=np.float32)
return ((raw - self.action_mean) / self.action_std).reshape(-1)
def __getitem__(self, idx):
ep = int(self.clip_episode[idx])
start = int(self.clip_start[idx])
base = int(self.offsets[ep])
length = int(self.lengths[ep])
# context frames, strided by frameskip
rows = [base + start + i * FRAMESKIP for i in range(NUM_CONTEXT)]
ctx = np.asarray(self.latents[rows], dtype=np.float32)
# blocks between the context frames: block i leaves context frame i
past = np.stack(
[
self._block(base, start + i * FRAMESKIP)
for i in range(NUM_CONTEXT - 1)
]
)
# goal: 1..max_offset transitions past the current frame, clamped to
# what is left in the episode
current = start + (NUM_CONTEXT - 1) * FRAMESKIP
reach = (length - 1 - current) // FRAMESKIP
high = min(self.max_offset, reach)
offset = int(torch.randint(1, high + 1, (1,)).item())
goal = np.asarray(
self.latents[base + current + offset * FRAMESKIP], dtype=np.float32
)
# the real block leaving the current frame — density model only
real = self._block(base, current)
return {
'context': torch.from_numpy(ctx),
'past_actions': torch.from_numpy(past),
'goal': torch.from_numpy(goal),
'real_action': torch.from_numpy(real),
'goal_offset': torch.tensor(offset, dtype=torch.long),
}
def split_episodes(num_episodes, val_fraction=0.05, seed=0):
"""Deterministic train/val split over episode indices."""
rng = np.random.default_rng(seed)
perm = rng.permutation(num_episodes)
n_val = max(1, int(num_episodes * val_fraction))
return perm[n_val:], perm[:n_val]
|