File size: 12,562 Bytes
20b0922 | 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 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | """Four-GPU DataParallel PAIR trainer.
This is the safe host-specific fallback when NCCL DDP is unavailable. It
still executes the frozen RGB-D policy's forward/backward replicas on all four
5090s; gradients are reduced by PyTorch to GPU0 rather than NCCL. The action
relation teacher is tiny and stays on GPU0 after the local policy outputs have
been gathered. Global B128 x 25k preserves the previous 3.2M-sample budget.
"""
from __future__ import annotations
import argparse
import glob
import json
import math
import random
from collections import OrderedDict, defaultdict
from pathlib import Path
import h5py
import numpy as np
import torch
import torch.nn.functional as F
from torch.nn import DataParallel
from torch.utils.data import DataLoader, Dataset, Sampler
from train_act import _stats, _trajectories, seed_everything
from train_stereo_act import DEPTH_MM_TO_M
from stereo_decoder_variants import PAIRActionTeacher, StereoPAIRAdapter
class CompactPAIRWristDataset(Dataset):
"""Stream-indexed RGB-D data without a Python object per control step.
The earlier dataset expands every frame of all 500 demonstrations into a
giant list before the first update. PAIR's synchronized sampler knows a
stream and time directly, so this compact representation eliminates that
multi-minute CPU/RAM bottleneck without changing any image/action value.
"""
def __init__(self, trajectories, horizon, stats, train, cache_limit=64):
self.horizon, self.stats, self.cache_limit = horizon, stats, cache_limit
kept = [item for index, item in enumerate(trajectories) if (index % 10 != 0) == train]
self.streams, self.episodes = [], []
for path, key, length, present, task in kept:
stream_ids = []
for arm in present:
stream_ids.append(len(self.streams))
self.streams.append((path, key, arm, task, length))
self.episodes.append((task, tuple(stream_ids), length))
self.cache = OrderedDict()
def __len__(self):
return sum(length * len(streams) for _task, streams, length in self.episodes)
def _episode(self, stream_id):
if stream_id not in self.cache:
path, key, arm, _task, _length = self.streams[stream_id]
with h5py.File(path, "r") as handle:
trajectory = handle[key]
sensor = trajectory["obs"]["sensor_data"][f"head_camera_agent{arm}"]
rgb, depth = sensor["rgb"][:], sensor["depth"][:]
if tuple(rgb.shape[1:]) != (480, 640, 3) or tuple(depth.shape[1:]) != (480, 640, 1):
raise ValueError(f"strict 640x480 RGB-D required for {path}:{key}:panda-{arm}")
self.cache[stream_id] = (rgb, depth,
trajectory["obs"]["agent"][f"panda-{arm}"]["qpos"][:].astype(np.float32),
trajectory["actions"][f"panda-{arm}"][:].astype(np.float32))
while len(self.cache) > self.cache_limit:
self.cache.popitem(last=False)
else:
self.cache.move_to_end(stream_id)
return self.cache[stream_id]
def __getitem__(self, request):
stream_id, time, group = request
rgb, depth, qpos, actions = self._episode(stream_id)
future = actions[time:time + self.horizon]
valid = len(future)
padded = np.empty((self.horizon, actions.shape[1]), np.float32)
padded[:valid], padded[valid:] = future, future[-1]
mask = np.zeros(self.horizon, np.bool_); mask[:valid] = True
return (torch.from_numpy(rgb[time]).permute(2, 0, 1).contiguous(),
torch.from_numpy(depth[time]).permute(2, 0, 1).contiguous(),
torch.from_numpy((qpos[time] - self.stats["q_mean"]) / self.stats["q_std"]),
torch.from_numpy((padded - self.stats["a_mean"]) / self.stats["a_std"]),
torch.from_numpy(mask), torch.tensor(group, dtype=torch.long))
class SameEpisodeTeamBlockSampler(Sampler):
"""One cached synchronized demonstration per 64-update block.
Each batch contains complete teams at many independently selected times.
This is both permutation invariant and I/O efficient: it keeps exactly
2/3/4 local RGB-D streams resident rather than repeatedly decoding scores
of long 640x480 demonstrations just to make one 120-sample batch.
"""
def __init__(self, dataset, batch_size, updates, block_updates, seed):
self.dataset, self.batch_size, self.updates = dataset, batch_size, updates
self.block_updates, self.seed, self.epoch = block_updates, seed, 0
self.by_task = defaultdict(list)
for episode in dataset.episodes:
self.by_task[episode[0]].append(episode)
self.tasks = sorted(self.by_task)
if len(self.tasks) != 5:
raise ValueError(f"expected all five tasks, got {self.tasks}")
def __len__(self): return self.updates
def __iter__(self):
rng = random.Random(self.seed + self.epoch); self.epoch += 1
done, block = 0, 0
while done < self.updates:
task = self.tasks[block % len(self.tasks)]
_task, streams, length = self.by_task[task][rng.randrange(len(self.by_task[task]))]
team_size = len(streams)
if self.batch_size % team_size:
raise ValueError("global batch must be divisible by every 2/3/4-agent team size")
for _ in range(min(self.block_updates, self.updates - done)):
batch = []
for group in range(self.batch_size // team_size):
time = rng.randrange(length)
batch.extend((stream, time, group) for stream in streams)
rng.shuffle(batch)
yield batch; done += 1
block += 1
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--data", required=True); parser.add_argument("--output", required=True)
parser.add_argument("--normalization", default=None,
help="existing same-corpus normalization.pt; avoids a redundant multi-GB scan")
parser.add_argument("--shared-arms", default="0,1,2,3")
parser.add_argument("--updates", type=int, default=26666)
parser.add_argument("--batch-size", type=int, default=120,
help="global batch divisible by 2/3/4-agent synchronized teams")
parser.add_argument("--episode-block-updates", type=int, default=64)
parser.add_argument("--cache-episodes", type=int, default=64)
parser.add_argument("--lr", type=float, default=3e-4); parser.add_argument("--warmup-updates", type=int, default=500)
parser.add_argument("--beta", type=float, default=1e-3); parser.add_argument("--roles", type=int, default=4)
parser.add_argument("--role-rank", type=int, default=32)
parser.add_argument("--distill-weight", type=float, default=.50)
parser.add_argument("--teacher-reconstruct-weight", type=float, default=.10)
parser.add_argument("--teacher-relation-weight", type=float, default=.10)
parser.add_argument("--teacher-usage-weight", type=float, default=.01)
parser.add_argument("--save-updates", default="26666"); parser.add_argument("--log-every", type=int, default=100)
parser.add_argument("--seed", type=int, default=20260730); parser.add_argument("--allow-preflight", action="store_true")
args = parser.parse_args()
if not args.allow_preflight and abs(args.batch_size * args.updates - 3_200_000) > args.batch_size:
raise ValueError("formal PAIR run must match the 3.2M-sample B40 x 80k budget within one batch")
if args.batch_size % 4:
raise ValueError("DataParallel batch must split equally across four GPUs")
seed_everything(args.seed); torch.backends.cudnn.benchmark = True
arms = tuple(int(value) for value in args.shared_arms.split(","))
paths = sorted({path for pattern in args.data.split(",") for path in glob.glob(pattern)})
trajectories = _trajectories(paths, arms)
if args.normalization:
stats = torch.load(args.normalization, map_location="cpu", weights_only=False)["stats"]
else:
stats = _stats(trajectories, arms)
dataset = CompactPAIRWristDataset(trajectories, 100, stats, True, cache_limit=args.cache_episodes)
sampler = SameEpisodeTeamBlockSampler(dataset, args.batch_size, args.updates, args.episode_block_updates, args.seed)
loader = DataLoader(dataset, batch_sampler=sampler, num_workers=0, pin_memory=True)
sample = dataset[(0, 0, 0)]; state_dim, action_dim = len(sample[2]), len(sample[3][0])
base = StereoPAIRAdapter(state_dim, action_dim, horizon=100, d_model=384, enc_layers=4,
dec_layers=7, roles=args.roles, role_rank=args.role_rank).cuda(0)
policy = DataParallel(base, device_ids=[0, 1, 2, 3], output_device=0)
teacher = PAIRActionTeacher(action_dim, roles=args.roles).cuda(0)
optimizer = torch.optim.AdamW(list(policy.parameters()) + list(teacher.parameters()), lr=args.lr, weight_decay=1e-4)
def schedule_multiplier(step):
warmup = min(1.0, (step + 1) / max(args.warmup_updates, 1))
return warmup * .5 * (1 + math.cos(math.pi * min(1.0, (step + 1) / args.updates)))
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, schedule_multiplier)
out = Path(args.output); out.mkdir(parents=True, exist_ok=True)
config = vars(args) | {"horizon":100,"enc_layers":4,"dec_layers":7,"d_model":384,
"vision_backbone":"stereo_act_cross_relbias","dino_model":"facebook/dinov3-vitb16-pretrain-lvd1689m",
"defm_model":base.defm_model_name,"camera_width":640,"camera_height":480,"patch_grid":[30,40],
"fusion_layers":2,"depth_storage_unit":"millimeters","depth_to_meters_scale":DEPTH_MM_TO_M,
"arms":arms,"state_dim":state_dim,"action_dim":action_dim,"files":paths,"episodes":len(trajectories),
"policy_variant":"stereo_pair_adapter","parallelism":"four-GPU PyTorch DataParallel; NCCL-free host fallback",
"global_batch":args.batch_size,"sample_budget":args.batch_size*args.updates,
"strict_policy_input":"current local panda_hand wrist RGB-D and local qpos only; no task/agent ID, peer/global/right-camera/language input",
"training_only_teacher":"permutation-invariant synchronized action-chunk relation teacher; absent at deployment"}
(out / "config.json").write_text(json.dumps(config, indent=2)); torch.save({"stats":stats},out / "normalization.pt")
milestones = {int(value) for value in args.save_updates.split(",") if value}; totals = {key:0.0 for key in ("loss","action","kl","distill","teacher_reconstruct","teacher_relation","teacher_usage")}
for update, batch in enumerate(loader, start=1):
rgb, depth, qpos, actions, mask, groups = [x.cuda(0, non_blocking=True) for x in batch]
optimizer.zero_grad(set_to_none=True)
with torch.autocast("cuda", dtype=torch.bfloat16):
prediction, mu, logvar, _aux, local_roles = policy(rgb.float().div_(255), depth, qpos, actions)
action_loss = (((prediction-actions).square().mean(-1)*mask).sum()/mask.sum().clamp_min(1))
kl = -.5*(1+logvar-mu.square()-logvar.exp()).sum(-1).mean()
teacher_roles, reconstruction, relation, usage = teacher(actions, groups)
distillation = F.kl_div(local_roles.clamp_min(1e-8).log(), teacher_roles.detach(), reduction="batchmean")
loss = action_loss + args.beta*kl + args.distill_weight*distillation + args.teacher_reconstruct_weight*reconstruction + args.teacher_relation_weight*relation + args.teacher_usage_weight*usage
loss.backward(); torch.nn.utils.clip_grad_norm_(list(policy.parameters())+list(teacher.parameters()),1.0)
optimizer.step(); scheduler.step()
values = {"loss":loss,"action":action_loss,"kl":kl,"distill":distillation,"teacher_reconstruct":reconstruction,"teacher_relation":relation,"teacher_usage":usage}
for name,value in values.items(): totals[name] += float(value.detach())
if update % args.log_every == 0 or update in milestones:
print(json.dumps({"update":update,"global_batch":args.batch_size,"lr":scheduler.get_last_lr()[0], **{key:value/update for key,value in totals.items()}}),flush=True)
if update in milestones:
torch.save({"model":policy.module.state_dict(),"stats":stats,"config":config,"update":update,"pair_teacher":teacher.state_dict()},out/f"checkpoint_{update:06d}.pt")
if __name__ == "__main__":
main()
|