Upload lewm_train.py

lewm_train.py

@@ -0,0 +1,404 @@
+"""
+LeWorldModel (LeWM) Training Script
+Reference: Maes et al., 2026 — Stable End-to-End JEPA from Pixels
+arXiv: 2603.19312
+
+This script trains LeWM on trajectory data (observations + actions).
+Supports both real HDF5 datasets and a synthetic PushT-like benchmark
+for rapid smoke-testing.
+"""
+
+import os
+import argparse
+import math
+import numpy as np
+import h5py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from einops import rearrange
+from transformers import get_cosine_schedule_with_warmup
+
+from lewm_model import build_lewm, SIGReg
+
+
+# ---------------------------------------------------------------------------
+# Dataset: HDF5 trajectory loader
+# ---------------------------------------------------------------------------
+class TrajectoryDataset(Dataset):
+    """
+    Loads offline trajectories from an HDF5 file.
+    Expected keys (standard from DINO-WM / LeWM datasets):
+        observations/pixels   (N_episodes, T_max, H, W, C) uint8
+        actions                 (N_episodes, T_max, A) float32
+    We extract sub-trajectories of length `seq_len` with frame_skip.
+    """
+    def __init__(self, h5_path, seq_len=4, frameskip=5, img_size=224,
+                 train=True, train_split=0.95):
+        self.seq_len = seq_len
+        self.frameskip = frameskip
+        self.img_size = img_size
+        self.train = train
+
+        with h5py.File(h5_path, 'r') as f:
+            pixels = f['observations']['pixels'][:]       # (N, T, H, W, C)
+            actions = f['actions'][:]                        # (N, T, A)
+
+        # Convert to torch tensors
+        self.pixels = torch.from_numpy(pixels).permute(0, 1, 4, 2, 3).float() / 255.0  # (N,T,C,H,W)
+        self.actions = torch.from_numpy(actions).float()
+
+        # Pre-compute episode boundaries
+        N, T_max = self.pixels.shape[:2]
+        n_train = int(N * train_split)
+        if train:
+            self.pixels = self.pixels[:n_train]
+            self.actions = self.actions[:n_train]
+        else:
+            self.pixels = self.pixels[n_train:]
+            self.actions = self.actions[n_train:]
+
+        N, T_max = self.pixels.shape[:2]
+        self.indices = []
+        for ep in range(N):
+            valid = T_max - (seq_len * frameskip) - 1
+            if valid > 0:
+                for start in range(0, valid, frameskip):
+                    self.indices.append((ep, start))
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, idx):
+        ep, start = self.indices[idx]
+        fs = self.frameskip
+        # Sample every frameskip-th frame
+        frame_indices = [start + t * fs for t in range(self.seq_len)]
+        obs = self.pixels[ep, frame_indices]           # (T, C, H, W)
+        # Actions: group `frameskip` consecutive actions into a block (mean or sum)
+        acts = []
+        for t in range(self.seq_len):
+            act_block = self.actions[ep, start + t * fs: start + (t + 1) * fs]
+            acts.append(act_block.mean(dim=0))
+        acts = torch.stack(acts, dim=0)                # (T, A)
+        return obs, acts
+
+
+# ---------------------------------------------------------------------------
+# Synthetic PushT-like dataset (for smoke-testing without 12 GB download)
+# ---------------------------------------------------------------------------
+class SyntheticPushTDataset(Dataset):
+    """
+    Generates synthetic 2D manipulation trajectories.
+    Agent (blue dot) pushes a T-shaped block toward a target.
+    Observations are rendered as 224×224 RGB images.
+    """
+    def __init__(self, n_episodes=2000, max_steps=196, img_size=224, seq_len=4, frameskip=5):
+        self.seq_len = seq_len
+        self.frameskip = frameskip
+        self.img_size = img_size
+        self.data = []
+        rng = np.random.RandomState(42)
+        min_steps = max(60, seq_len * frameskip + 10)
+        for _ in range(n_episodes):
+            length = rng.randint(min_steps, max( min_steps + 1, max_steps))
+            traj = self._generate_trajectory(length, rng)
+            self.data.append(traj)
+
+    def _generate_trajectory(self, length, rng):
+        img_size = self.img_size
+        # Agent pos, block pos, block angle
+        agent = rng.uniform(0.2, 0.8, size=(length, 2)).astype(np.float32)
+        block = rng.uniform(0.3, 0.7, size=(length, 2)).astype(np.float32)
+        angle = np.cumsum(rng.randn(length).astype(np.float32) * 0.1)
+        # Actions: dx, dy for agent (2D continuous)
+        actions = np.diff(agent, prepend=agent[:1], axis=0).astype(np.float32)
+        # Pad to uniform length by repeating last frame
+        pixels = np.zeros((length, 3, img_size, img_size), dtype=np.float32)
+        for t in range(length):
+            pixels[t] = self._render(agent[t], block[t], angle[t], img_size)
+        return {"pixels": pixels, "actions": actions}
+
+    @staticmethod
+    def _render(agent, block, angle, size):
+        canvas = np.ones((3, size, size), dtype=np.float32) * 0.9
+        # Draw agent (blue circle)
+        y, x = np.ogrid[:size, :size]
+        ax, ay = int(agent[0] * size), int(agent[1] * size)
+        mask = ((x - ax) ** 2 + (y - ay) ** 2) < (size * 0.03) ** 2
+        canvas[2][mask] = 0.3
+        canvas[0][mask] = 0.3
+        # Draw block (red T)
+        bx, by = int(block[0] * size), int(block[1] * size)
+        block_mask = ((x - bx) ** 2 + (y - by) ** 2) < (size * 0.05) ** 2
+        canvas[0][block_mask] = 0.9
+        canvas[1][block_mask] = 0.2
+        canvas[2][block_mask] = 0.2
+        return canvas
+
+    def __len__(self):
+        return len(self.data) * 50  # many sub-trajectories per episode
+
+    def __getitem__(self, idx):
+        ep = idx % len(self.data)
+        traj = self.data[ep]
+        max_start = len(traj["pixels"]) - self.seq_len * self.frameskip - 1
+        if max_start <= 0:
+            max_start = 1
+        start = np.random.randint(0, max_start)
+        fs = self.frameskip
+        frame_idx = [start + t * fs for t in range(self.seq_len)]
+        obs = torch.from_numpy(traj["pixels"][frame_idx])
+        acts = []
+        for t in range(self.seq_len):
+            a = traj["actions"][start + t * fs: start + (t + 1) * fs].mean(axis=0)
+            acts.append(a)
+        acts = torch.from_numpy(np.stack(acts, axis=0))
+        # Pad actions to effective dim (frameskip * action_dim)
+        A = acts.shape[-1]
+        pad = fs * A - A
+        if pad > 0:
+            acts = F.pad(acts, (0, pad))
+        return obs, acts
+
+
+# ---------------------------------------------------------------------------
+# Training loop
+# ---------------------------------------------------------------------------
+def train(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    # Build model
+    model = build_lewm(
+        image_size=args.img_size,
+        patch_size=14,
+        embed_dim=args.embed_dim,
+        action_dim=args.action_dim,
+        history_size=args.history_size,
+        frameskip=args.frameskip,
+        predictor_depth=6,
+        predictor_heads=16,
+        predictor_mlp_dim=2048,
+        predictor_dropout=0.1,
+    ).to(device)
+
+    print(f"Model params: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")
+
+    # Dataset
+    if args.use_synthetic:
+        dataset = SyntheticPushTDataset(
+            n_episodes=args.n_episodes,
+            seq_len=args.seq_len,
+            frameskip=args.frameskip,
+            img_size=args.img_size,
+        )
+        val_dataset = SyntheticPushTDataset(
+            n_episodes=max(1, args.n_episodes // 10),
+            seq_len=args.seq_len,
+            frameskip=args.frameskip,
+            img_size=args.img_size,
+        )
+    else:
+        dataset = TrajectoryDataset(
+            args.h5_path, seq_len=args.seq_len, frameskip=args.frameskip,
+            img_size=args.img_size, train=True,
+        )
+        val_dataset = TrajectoryDataset(
+            args.h5_path, seq_len=args.seq_len, frameskip=args.frameskip,
+            img_size=args.img_size, train=False,
+        )
+
+    loader = DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=True,
+        num_workers=args.num_workers, drop_last=True, pin_memory=True,
+    )
+    val_loader = DataLoader(
+        val_dataset, batch_size=args.batch_size, shuffle=False,
+        num_workers=0, drop_last=False, pin_memory=True,
+    )
+
+    # Optimizer + scheduler
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay,
+        betas=(0.9, 0.95),
+    )
+    total_steps = len(loader) * args.epochs
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer, num_warmup_steps=int(0.05 * total_steps),
+        num_training_steps=total_steps,
+    )
+
+    # SIGReg
+    sigreg = SIGReg(knots=17, num_proj=1024).to(device)
+
+    # Training
+    best_val_loss = float('inf')
+    for epoch in range(args.epochs):
+        model.train()
+        epoch_loss = 0.0
+        epoch_pred = 0.0
+        epoch_sig = 0.0
+
+        for step, (obs, acts) in enumerate(loader):
+            obs = obs.to(device)
+            acts = acts.to(device)
+            b, t = obs.shape[:2]
+
+            # Encode
+            emb = model.encode(obs)  # (B, T, D)
+            act_emb = model.action_encoder(acts)
+
+            # Predictor (history_size)
+            ctx_emb = emb[:, :args.history_size]
+            ctx_act = act_emb[:, :args.history_size]
+            pred_emb = model.predict(ctx_emb, ctx_act)
+
+            # Prediction loss
+            pred_loss = (pred_emb[:, :-1] - emb[:, 1:args.history_size]).pow(2).mean()
+
+            # SIGReg
+            sigreg_loss = sigreg(emb.transpose(0, 1))
+
+            loss = pred_loss + args.lambd * sigreg_loss
+
+            optimizer.zero_grad()
+            loss.backward()
+            if args.grad_clip > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
+            optimizer.step()
+            scheduler.step()
+
+            epoch_loss += loss.item()
+            epoch_pred += pred_loss.item()
+            epoch_sig += sigreg_loss.item()
+
+            if step % args.log_interval == 0:
+                print(f"  [E{epoch}|S{step}] loss={loss.item():.4f} "
+                      f"pred={pred_loss.item():.4f} sigreg={sigreg_loss.item():.4f}")
+
+        n = len(loader)
+        print(f"Epoch {epoch} | avg_loss={epoch_loss/n:.4f} "
+              f"avg_pred={epoch_pred/n:.4f} avg_sigreg={epoch_sig/n:.4f}")
+
+        # Validation
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for obs, acts in val_loader:
+                obs = obs.to(device)
+                acts = acts.to(device)
+                emb = model.encode(obs)
+                act_emb = model.action_encoder(acts)
+                ctx_emb = emb[:, :args.history_size]
+                ctx_act = act_emb[:, :args.history_size]
+                pred_emb = model.predict(ctx_emb, ctx_act)
+                pred_loss = (pred_emb[:, :-1] - emb[:, 1:args.history_size]).pow(2).mean()
+                sigreg_loss = sigreg(emb.transpose(0, 1))
+                val_loss += (pred_loss + args.lambd * sigreg_loss).item()
+        val_loss /= max(1, len(val_loader))
+        print(f"  Val loss: {val_loss:.4f}")
+
+        # Save best
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            ckpt = {
+                "model": model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "scheduler": scheduler.state_dict(),
+                "epoch": epoch,
+                "args": vars(args),
+            }
+            out_path = os.path.join(args.output_dir, "best_model.pt")
+            os.makedirs(args.output_dir, exist_ok=True)
+            torch.save(ckpt, out_path)
+            print(f"  Saved best model -> {out_path}")
+
+    # Final save
+    final_path = os.path.join(args.output_dir, "final_model.pt")
+    torch.save({"model": model.state_dict(), "args": vars(args)}, final_path)
+    print(f"Training complete. Saved to {final_path}")
+
+    # Push to hub
+    if args.push_to_hub:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        repo_id = f"{args.hf_username}/{args.hub_model_id}"
+        api.create_repo(repo_id, repo_type="model", exist_ok=True)
+        api.upload_file(
+            path_or_fileobj=final_path,
+            path_in_repo="model.pt",
+            repo_id=repo_id,
+            repo_type="model",
+        )
+        # Save config
+        import json
+        config = {
+            "_target_": "lewm_model.LeWorldModel",
+            "encoder": {
+                "image_size": args.img_size,
+                "patch_size": 14,
+                "embed_dim": args.embed_dim,
+                "num_layers": 12,
+                "num_heads": 3,
+            },
+            "predictor": {
+                "num_frames": args.history_size,
+                "depth": 6,
+                "heads": 16,
+                "mlp_dim": 2048,
+                "dropout": 0.1,
+            },
+            "action_dim": args.action_dim,
+            "frameskip": args.frameskip,
+            "lambd": args.lambd,
+        }
+        config_path = os.path.join(args.output_dir, "config.json")
+        with open(config_path, "w") as f:
+            json.dump(config, f, indent=2)
+        api.upload_file(
+            path_or_fileobj=config_path,
+            path_in_repo="config.json",
+            repo_id=repo_id,
+            repo_type="model",
+        )
+        print(f"Pushed model to https://huggingface.co/{repo_id}")
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+def get_args():
+    parser = argparse.ArgumentParser(description="Train LeWorldModel")
+    # Data
+    parser.add_argument("--h5_path", type=str, default="/tmp/pusht_expert_train.h5")
+    parser.add_argument("--use_synthetic", action="store_true", help="Use synthetic data for smoke testing")
+    parser.add_argument("--n_episodes", type=int, default=2000, help="Synthetic dataset size")
+    parser.add_argument("--seq_len", type=int, default=4)
+    parser.add_argument("--frameskip", type=int, default=5)
+    parser.add_argument("--img_size", type=int, default=224)
+    parser.add_argument("--action_dim", type=int, default=2)
+    parser.add_argument("--history_size", type=int, default=3)
+    # Model
+    parser.add_argument("--embed_dim", type=int, default=192)
+    parser.add_argument("--lambd", type=float, default=0.1, help="SIGReg weight")
+    # Training
+    parser.add_argument("--epochs", type=int, default=10)
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--weight_decay", type=float, default=0.05)
+    parser.add_argument("--grad_clip", type=float, default=1.0)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--log_interval", type=int, default=50)
+    parser.add_argument("--output_dir", type=str, default="/tmp/lewm_output")
+    # Hub
+    parser.add_argument("--push_to_hub", action="store_true")
+    parser.add_argument("--hf_username", type=str, default="ar27111994")
+    parser.add_argument("--hub_model_id", type=str, default="lewm-synthetic-pusht")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = get_args()
+    train(args)