train.py

import math
import os
import sys
import time
from ast import literal_eval
from collections import deque
from pathlib import Path

import gymnasium as gym
import numpy as np
import torch
import torch.distributed as dist
import torch.nn.functional as F
from gymnasium.envs.box2d.lunar_lander import heuristic
from PIL import Image
from torch import nn
from torch.nn.parallel import DistributedDataParallel as DDP

os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")

# model
image_size = 64
history = 4
patch_size = 4
channels = 3
obs_dim = 8
width = 768
depth = 8
heads = 12

# data
replay_capacity = 4096
prefill_windows = 64
gravity = -10.0
wind_power = 15.0
turbulence_power = 2.0

# training
steps = 5_500
batch_size = 64
micro_batch_size = 32
learning_rate = 1e-3
seed = 0
log_every = 50
save_every = 500
video_every = 500
video_episodes = 4
denoising_steps = 4
checkpoint = "checkpoint.pt"
compile_model = True
wandb_log = False

config_keys = [
    key
    for key, value in globals().items()
    if not key.startswith("_") and isinstance(value, (int, float, bool, str))
]
for arg in sys.argv[1:]:
    key, value = arg.removeprefix("--").split("=", 1)
    try:
        value = literal_eval(value)
    except (SyntaxError, ValueError):
        pass
    assert key in config_keys and type(value) is type(globals()[key])
    globals()[key] = value
config = {key: globals()[key] for key in config_keys}
num_frames = history + 1


def pixels(frame):
    image = Image.fromarray(frame).resize(
        (image_size, image_size), Image.Resampling.BILINEAR
    )
    return (
        torch.from_numpy(np.asarray(image, dtype=np.uint8).copy())
        .permute(2, 0, 1)
        .float()
        .div_(127.5)
        .sub_(1)
    )


def sincos(position, width):
    frequency = 1 / 10_000 ** (
        torch.arange(width // 2, dtype=torch.float32) / (width / 2)
    )
    phase = position.float().flatten()[:, None] * frequency
    return torch.cat((phase.sin(), phase.cos()), dim=1)


def make_env():
    return gym.make(
        "LunarLander-v3",
        continuous=True,
        render_mode="rgb_array",
        gravity=gravity,
        enable_wind=True,
        wind_power=wind_power,
        turbulence_power=turbulence_power,
    )


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm1 = nn.RMSNorm(width, eps=1e-5, elementwise_affine=False)
        self.qkv = nn.Linear(width, 3 * width)
        self.proj = nn.Linear(width, width)
        self.norm2 = nn.RMSNorm(width, eps=1e-5, elementwise_affine=False)
        self.mlp = nn.Sequential(
            nn.Linear(width, 4 * width),
            nn.GELU(approximate="tanh"),
            nn.Linear(4 * width, width),
        )
        self.modulation = nn.Sequential(nn.SiLU(), nn.Linear(width, 6 * width))

    def forward(self, x, condition, attention_mask):
        shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = self.modulation(
            condition
        ).chunk(6, dim=-1)
        q, k, v = (
            self.qkv(self.norm1(x) * (1 + scale_a) + shift_a)
            .view(*x.shape[:2], 3, heads, -1)
            .unbind(2)
        )
        attention = F.scaled_dot_product_attention(
            q.transpose(1, 2),
            k.transpose(1, 2),
            v.transpose(1, 2),
            attn_mask=attention_mask,
        )
        x = x + gate_a * self.proj(attention.transpose(1, 2).flatten(2))
        return x + gate_m * self.mlp(self.norm2(x) * (1 + scale_m) + shift_m)


class WorldModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.patch = nn.Conv3d(
            channels,
            width,
            (1, patch_size, patch_size),
            stride=(1, patch_size, patch_size),
        )
        self.time = nn.Sequential(
            nn.Linear(128, width), nn.SiLU(), nn.Linear(width, width)
        )
        self.observation = nn.Sequential(
            nn.Linear(obs_dim, width), nn.SiLU(), nn.Linear(width, width)
        )
        self.blocks = nn.ModuleList(Block() for _ in range(depth))
        self.norm = nn.RMSNorm(width, eps=1e-5, elementwise_affine=False)
        self.out = nn.Linear(width, channels * patch_size * patch_size)

        grid = image_size // patch_size
        y, x = torch.meshgrid(torch.arange(grid), torch.arange(grid), indexing="ij")
        spatial = torch.cat((sincos(x, width // 2), sincos(y, width // 2)), dim=1)
        position = spatial.repeat(num_frames, 1) + sincos(
            torch.arange(num_frames), width
        ).repeat_interleave(grid * grid, dim=0)
        frame = torch.arange(num_frames).repeat_interleave(grid * grid)
        self.register_buffer("position", position[None], persistent=False)
        self.register_buffer(
            "attention_mask",
            (frame[:, None] >= frame[None, :])[None, None],
            persistent=False,
        )
        self.register_buffer(
            "frequency",
            torch.exp(-math.log(10_000) * torch.arange(64) / 64),
            persistent=False,
        )

    def forward(self, frames, timestep, observation):
        batch = frames.shape[0]
        grid = image_size // patch_size
        # [B, T, C, H, W]
        x = self.patch(frames.transpose(1, 2)).flatten(2).transpose(1, 2)
        x = x + self.position.to(x.dtype)
        phase = timestep.float()[..., None] * self.frequency
        condition = self.time(
            torch.cat((phase.cos(), phase.sin()), dim=-1)
        ) + self.observation(observation)
        condition = condition.repeat_interleave(grid * grid, dim=1)
        for block in self.blocks:
            x = block(x, condition, self.attention_mask)
        x = self.out(self.norm(x)).view(
            batch,
            num_frames,
            grid,
            grid,
            channels,
            patch_size,
            patch_size,
        )
        return x.permute(0, 1, 4, 2, 5, 3, 6).reshape(
            batch, num_frames, channels, image_size, image_size
        )


class Replay:
    def __init__(self, rank, seed):
        self.seed = seed + 100_000 * rank
        self.rng = np.random.default_rng(self.seed)
        self.generator = torch.Generator().manual_seed(self.seed)
        self.buffer = deque(maxlen=replay_capacity)
        self.frames = deque(maxlen=num_frames)
        self.observations = deque(maxlen=num_frames)
        self.episodes = 0
        self.env = make_env()
        self.reset()

    def reset(self):
        self.observation, _ = self.env.reset(seed=self.seed + self.episodes)
        self.episodes += 1
        self.frames.clear()
        self.observations.clear()
        self.frames.append(pixels(self.env.render()))
        self.observations.append(torch.from_numpy(self.observation.copy()))

    def step(self):
        action = np.asarray(
            heuristic(self.env.unwrapped, self.observation), dtype=np.float32
        )
        action = np.clip(
            action, self.env.action_space.low, self.env.action_space.high
        ).astype(np.float32)
        self.observation, _, terminated, truncated, _ = self.env.step(action)
        self.frames.append(pixels(self.env.render()))
        self.observations.append(torch.from_numpy(self.observation.copy()))
        if len(self.frames) == num_frames:
            self.buffer.append(
                (torch.stack(tuple(self.frames)), torch.stack(tuple(self.observations)))
            )
        if terminated or truncated:
            self.reset()

    def batch(self, size):
        while len(self.buffer) < prefill_windows:
            self.step()
        samples = []
        for _ in range(size):
            self.step()
            samples.append(self.buffer[int(self.rng.integers(len(self.buffer)))])
        frames = torch.stack([sample[0] for sample in samples])
        observations = torch.stack([sample[1] for sample in samples])

        noise = torch.randn(frames.shape, generator=self.generator)
        noise_t = torch.sigmoid(
            torch.randn((size, num_frames), generator=self.generator)
        )
        num_noisy_history = torch.randint(
            history + 1, (size, 1), generator=self.generator
        )
        noisy_history = torch.arange(history)[None] >= history - num_noisy_history
        mask = torch.cat((noisy_history, torch.ones((size, 1), dtype=torch.bool)), 1)
        noise_t.masked_fill_(~mask, 0)
        noised = (1 - noise_t[..., None, None, None]) * frames + noise_t[
            ..., None, None, None
        ] * noise
        model_t = noise_t.clone()
        model_t[:, :-1].masked_fill_(noisy_history, 0)
        return (
            noised,
            model_t,
            observations,
            frames - noise,
            mask[..., None, None, None],
        )

    def close(self):
        self.env.close()


def to_rgb(frame):
    return frame.clamp(-1, 1).add(1).mul(127.5).byte().permute(1, 2, 0).numpy()


@torch.no_grad()
def rollout(model, device, seed, denoising_steps):
    env = make_env()
    generator = torch.Generator(device=device).manual_seed(seed)
    observation, _ = env.reset(seed=seed)
    real = pixels(env.render())
    context = deque([real], maxlen=history)
    observations = deque([torch.from_numpy(observation.copy())], maxlen=history)
    video = [np.concatenate((to_rgb(real), to_rgb(real)), axis=1)]
    terminated = truncated = False
    while not (terminated or truncated):
        action = np.asarray(heuristic(env.unwrapped, observation), dtype=np.float32)
        action = np.clip(action, env.action_space.low, env.action_space.high).astype(
            np.float32
        )
        observation, _, terminated, truncated, _ = env.step(action)
        real = pixels(env.render())
        if len(context) < history:
            predicted = real
        else:
            context_frames = torch.stack(tuple(context)).to(device)
            state = torch.stack(
                (*observations, torch.from_numpy(observation.copy()))
            ).to(device)[None]
            predicted = torch.randn(
                (channels, image_size, image_size), generator=generator, device=device
            )
            timestep = torch.zeros((1, num_frames), device=device)
            # TODO: add CFG, doesn't seem to need it
            for i in range(denoising_steps):
                timestep[:, -1] = 1 - i / denoising_steps
                frames = torch.cat((context_frames, predicted[None]))[None]
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    velocity = model(frames, timestep, state)
                predicted.add_(velocity[0, -1].float(), alpha=1 / denoising_steps)
            predicted = predicted.nan_to_num().clamp(-1, 1).cpu()
        context.append(predicted)
        observations.append(torch.from_numpy(observation.copy()))
        video.append(np.concatenate((to_rgb(predicted), to_rgb(real)), axis=1))
    env.close()
    return np.stack(video)


def main():
    rank = int(os.environ.get("RANK", "0"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    world = int(os.environ.get("WORLD_SIZE", "1"))
    ddp = world > 1
    master = rank == 0
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    if ddp:
        dist.init_process_group("nccl", device_id=device)
    torch.manual_seed(seed)

    raw_model = WorldModel().to(device)
    if master:
        print(
            f"model: {sum(p.numel() for p in raw_model.parameters()) / 1e6:.2f}M parameters",
            flush=True,
        )
    model = torch.compile(raw_model) if compile_model else raw_model
    if ddp:
        model = DDP(model, device_ids=[local_rank])
    optimizer = torch.optim.AdamW(
        raw_model.parameters(),
        lr=learning_rate,
        betas=(0.9, 0.95),
        weight_decay=0.01,
        fused=True,
    )

    warmup = max(1, steps // 100)
    decay = round(steps * 0.1)
    stable = steps + 1 - warmup - decay

    def get_lr(step):
        if step <= warmup:
            return learning_rate * step / warmup
        if step <= warmup + stable:
            return learning_rate
        progress = (step - warmup - stable) / decay
        return learning_rate * 0.5 * (1 + math.cos(math.pi * progress))

    start = 0
    checkpoint_path = Path(checkpoint)
    if checkpoint_path.exists():
        state = torch.load(checkpoint_path, map_location=device)
        raw_model.load_state_dict(state["model"])
        optimizer.load_state_dict(state["optimizer"])
        start = state["step"]

    replay = Replay(rank, seed)
    run = None
    if master and wandb_log:
        import wandb

        run = wandb.init(
            project=os.getenv("WANDB_PROJECT", "nano-worldmodel"),
            config=config,
            resume="allow",
        )

    for step in range(start + 1, steps + 1):
        tick = time.perf_counter()
        batch = replay.batch(batch_size)
        valid = (
            batch[-1].sum().to(device=device, dtype=torch.float32)
            * channels
            * image_size
            * image_size
        )
        if ddp:
            dist.all_reduce(valid)

        lr = get_lr(step)
        for group in optimizer.param_groups:
            group["lr"] = lr
        optimizer.zero_grad(set_to_none=True)
        total_error = torch.zeros((), device=device)
        for offset in range(0, batch_size, micro_batch_size):
            end = min(offset + micro_batch_size, batch_size)
            if ddp:
                model.require_backward_grad_sync = end == batch_size
            frames, timestep, observation, target, mask = (
                x[offset:end].to(device) for x in batch
            )
            with torch.autocast("cuda", dtype=torch.bfloat16):
                prediction = model(frames, timestep, observation)
            squared_error = ((prediction.float() - target) ** 2 * mask).sum()
            loss = squared_error * world / valid
            loss.backward()
            total_error += squared_error.detach()
        grad_norm = nn.utils.clip_grad_norm_(raw_model.parameters(), 1.0)
        optimizer.step()

        logging = step == 1 or step % log_every == 0
        if logging:
            episodes = torch.tensor(replay.episodes, device=device)
            if ddp:
                dist.all_reduce(total_error)
                dist.all_reduce(episodes)
            metrics = {
                "loss": (total_error / valid).item(),
                "grad_norm": grad_norm.item(),
                "lr": lr,
                "episodes": episodes.item(),
                "step_seconds": time.perf_counter() - tick,
            }
            if master:
                print(
                    f"step {step:5d} | loss {metrics['loss']:.5f} | grad {metrics['grad_norm']:.3f} | {metrics['step_seconds']:.2f}s",
                    flush=True,
                )
                if run:
                    run.log(metrics, step=step)

        saving = step == 1 or step % save_every == 0 or step == steps
        if master and saving:
            checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
            torch.save(
                {
                    "model": raw_model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "step": step,
                },
                checkpoint_path,
            )
        if ddp and saving:
            dist.barrier()

        filming = wandb_log and (step == 1 or step % video_every == 0)
        if filming:
            if ddp:
                dist.barrier()
            if master:
                raw_model.eval()
                videos = {}
                for episode in range(video_episodes):
                    frames = rollout(
                        raw_model, device, step * 10_000 + episode, denoising_steps
                    )
                    videos[f"validation_videos/nano_worldmodel_episode_{episode}"] = (
                        wandb.Video(
                            np.moveaxis(frames, -1, 1),
                            caption=f"nano worldmodel validation episode {episode}",
                            fps=50,
                            format="mp4",
                        )
                    )
                run.log(videos, step=step)
                raw_model.train()
            if ddp:
                dist.barrier()

    replay.close()
    if run:
        run.finish()
    if ddp:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()