File size: 5,092 Bytes
68eed61 | 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 | #!/usr/bin/env python3
"""Small real-data training smoke for bridge_orig_lerobot_smoke.
This deliberately tests decoding, batching, CUDA forward/backward, optimizer
updates, and checkpoint writing without claiming to train the full Cosmos/Zeva
model. The Zeva framework's Bridge adapter currently expects older LeRobot
metadata, while this smoke set is v3.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
import pyarrow.parquet as pq
import torch
from torch import nn
import av
class TinyActionNet(nn.Module):
def __init__(self) -> None:
super().__init__()
self.encoder = nn.Sequential(
nn.Conv2d(3, 16, 5, stride=2, padding=2), nn.GELU(),
nn.Conv2d(16, 32, 5, stride=2, padding=2), nn.GELU(),
nn.Conv2d(32, 64, 5, stride=2, padding=2), nn.GELU(),
nn.AdaptiveAvgPool2d(1),
)
self.head = nn.Sequential(nn.Linear(64 + 8, 64), nn.GELU(), nn.Linear(64, 7))
def forward(self, image: torch.Tensor, state: torch.Tensor) -> torch.Tensor:
feat = self.encoder(image).flatten(1)
return self.head(torch.cat([feat, state], dim=1))
def load_subset(root: Path, episodes: int, size: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
rows_i, images, states, actions = [], [], [], []
for episode in range(episodes):
table = pq.read_table(root / "data" / "chunk-000" / f"episode_{episode:06d}.parquet").to_pylist()
video = root / "videos" / "chunk-000" / "observation.images.image_0" / f"episode_{episode:06d}.mp4"
container = av.open(str(video))
decoded = [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)]
container.close()
if len(decoded) < len(table):
raise RuntimeError(f"decoded {len(decoded)} frames but parquet has {len(table)}: {video}")
for row, frame in zip(table, decoded):
# Avoid an OpenCV dependency; torch interpolation handles the resize.
frame = torch.from_numpy(frame).permute(2, 0, 1).float().unsqueeze(0)
frame = torch.nn.functional.interpolate(frame, size=(size, size), mode="bilinear", align_corners=False)
frame = frame.squeeze(0).byte().permute(1, 2, 0).numpy()
images.append(frame)
states.append(row["observation.state"])
actions.append(row["action"])
if not images:
raise RuntimeError("smoke subset is empty")
return (
torch.from_numpy(np.asarray(images)).permute(0, 3, 1, 2).float().div_(255.0),
torch.tensor(np.asarray(states), dtype=torch.float32),
torch.tensor(np.asarray(actions), dtype=torch.float32),
)
def main() -> int:
p = argparse.ArgumentParser()
p.add_argument("--root", type=Path, required=True)
p.add_argument("--steps", type=int, default=2000)
p.add_argument("--episodes", type=int, default=3)
p.add_argument("--output", type=Path, required=True)
p.add_argument("--device", default="cuda")
args = p.parse_args()
if args.steps < 1:
raise ValueError("steps must be positive")
device = torch.device(args.device if args.device != "cuda" or torch.cuda.is_available() else "cpu")
images, states, actions = load_subset(args.root, args.episodes, 64)
images, states, actions = images.to(device), states.to(device), actions.to(device)
model = TinyActionNet().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
loss_fn = nn.MSELoss()
losses: list[float] = []
generator = torch.Generator(device=device).manual_seed(42)
model.train()
for step in range(args.steps):
idx = torch.randint(images.shape[0], (min(8, images.shape[0]),), generator=generator, device=device)
loss = loss_fn(model(images.index_select(0, idx), states.index_select(0, idx)), actions.index_select(0, idx))
opt.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
losses.append(float(loss.detach().cpu()))
if (step + 1) % max(1, args.steps // 10) == 0:
print(f"step={step + 1} loss={losses[-1]:.8f}", flush=True)
args.output.parent.mkdir(parents=True, exist_ok=True)
torch.save({"model": model.state_dict(), "steps": args.steps}, args.output.with_suffix(".pt"))
report = {
"status": "PASS_DATA_GPU_TRAINING_SMOKE",
"device": str(device),
"episodes": args.episodes,
"frames": int(images.shape[0]),
"image_shape": list(images.shape[1:]),
"state_shape": list(states.shape[1:]),
"action_shape": list(actions.shape[1:]),
"steps": args.steps,
"loss_first": losses[0],
"loss_last": losses[-1],
"loss_min": min(losses),
"checkpoint": str(args.output.with_suffix(".pt").resolve()),
}
args.output.write_text(json.dumps(report, indent=2), encoding="utf-8")
print(json.dumps(report, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())
|