Project: build a mini driving VLA

def teacher(
    gap: torch.Tensor, speed: torch.Tensor, lead: torch.Tensor, style: torch.Tensor
) -> torch.Tensor:
    """A time-gap controller whose following distance is set by the instruction.

    It aims for 3 m plus ``HEADWAY[style]`` seconds of travel, matches the lead's
    speed, and brakes at full strength when the time to collision falls under 2 s.
    """
    wanted = 3.0 + HEADWAY[style] * speed
    accel = 0.4 * (gap - wanted) + 1.0 * (lead - speed)
    closing = (speed - lead).clamp(min=1e-3)
    emergency = (gap / closing < 2.0) & (speed > lead)
    accel = torch.where(emergency, torch.full_like(accel, -MAX_BRAKE), accel)
    return accel.clamp(-MAX_BRAKE, MAX_ACCEL)

def to_token(accel: torch.Tensor) -> torch.Tensor:
    """The index of the nearest action bin: an acceleration becomes a token."""
    return (accel[..., None] - BINS).abs().argmin(dim=-1)

class MiniDriver(nn.Module):
    """Instruction token + observation tokens -> transformer -> action-token logits.

    The instruction is one learned embedding (like a word); each of the last
    ``HISTORY`` observations (gap, speed, closing speed) is projected to a token and
    given a learned position. The output at the instruction token is read out as
    logits over the ``BINS`` action vocabulary.
    """

    def __init__(self, dim: int = 32, heads: int = 4, layers: int = 2) -> None:
        """Creates the embeddings, the transformer blocks and the action head."""
        super().__init__()
        self.embed = nn.Linear(3, dim)
        self.position = nn.Parameter(0.02 * torch.randn(HISTORY, dim))
        self.instruction = nn.Embedding(len(HEADWAY), dim)
        self.blocks = nn.ModuleList(TransformerBlock(dim, heads) for _ in range(layers))
        self.head = nn.Linear(dim, len(BINS))

    def forward(self, history: torch.Tensor, style: torch.Tensor) -> torch.Tensor:
        """Maps history (batch, HISTORY, 3) and style (batch,) to logits."""
        tokens = torch.cat(
            [self.instruction(style)[:, None], self.embed(history) + self.position],
            dim=1,
        )
        for block in self.blocks:
            tokens = block(tokens)
        return self.head(tokens[:, 0])

def demonstrations(
    scenes: Scenes, style: torch.Tensor, policy: Policy = teacher_policy
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Drives with ``policy``; returns each visited (history, style, teacher token)."""
    result = drive(scenes, style, policy, record=True)
    n = len(style)
    return (
        result["history"].reshape(n * STEPS, HISTORY, 3),
        style[:, None].expand(n, STEPS).reshape(-1),
        to_token(result["teacher"]).reshape(-1),
    )


def fit(
    model: MiniDriver,
    data: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
    steps: int = 600,
    batch: int = 2048,
) -> float:
    """Behaviour cloning: cross-entropy of the model's logits on teacher tokens."""
    history, style, target = data
    optimizer = torch.optim.Adam(model.parameters(), lr=3e-3)
    generator = torch.Generator().manual_seed(0)
    for _ in range(steps):
        index = torch.randint(0, len(history), (batch,), generator=generator)
        loss = nn.functional.cross_entropy(
            model(history[index], style[index]), target[index]
        )
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return loss.item()

def data_engine(
    data: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
    model: MiniDriver,
    rounds: int = 2,
    pool: int = 2000,
    added: int = 8,
    seed: int = 0,
) -> tuple[MiniDriver, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
    """Drives a fresh pool, keeps scenarios that crashed, has the teacher label them.

    Each round the model drives ``pool`` new scenarios (2 percent hard-braking
    leads); up to ``added`` of those it crashed in are driven again by the model with
    the teacher labelling every visited state, and a fresh model is fitted on all data.
    Stops early when the pool has no crash.
    """
    for round_ in range(rounds):
        scenes = make_scenarios(pool, 0.02, seed=500 + 10 * seed + round_)
        style = instructions(pool, 600 + 10 * seed + round_)
        crashed = drive(scenes, style, greedy(model))["crashed"]
        pick = crashed.nonzero().squeeze(-1)[:added]
        if len(pick) == 0:
            break
        failing = {key: value[pick] for key, value in scenes.items()}
        new = demonstrations(failing, style[pick], greedy(model))
        data = tuple(
            torch.cat([old, extra]) for old, extra in zip(data, new, strict=True)
        )
        torch.manual_seed(seed)
        model = MiniDriver()
        fit(model, data)
    return model, data

def episode_reward(
    result: dict[str, torch.Tensor],
    collision_weight: float = 5.0,
    comfort_weight: float = 0.05,
) -> torch.Tensor:
    """Progress (mean speed / 10) minus a collision penalty and a jerk penalty."""
    progress = result["speed"].mean(dim=1) / 10.0
    jerk = result["accel"].diff(dim=1).abs().mean(dim=1) / DT
    return (
        progress - collision_weight * result["crashed"].float() - comfort_weight * jerk
    )

def grpo_finetune(
    model: MiniDriver,
    iterations: int = 40,
    prompts: int = 32,
    group: int = 8,
    lr: float = 3e-4,
    kl_coef: float = 0.04,
    reward: Callable[[dict[str, torch.Tensor]], torch.Tensor] = episode_reward,
    seed: int = 0,
) -> list[float]:
    """GRPO: each scenario is a prompt, ``group`` sampled drives are its outputs.

    Every drive gets one reward; its advantage, against the other drives of the same
    scenario, is given to every action token it took. The loss is the clipped ratio of
    each token's probability with a k3 KL penalty to the starting model.

    Returns:
        The mean training reward at each iteration.
    """
    reference = copy.deepcopy(model).eval()
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    history_of_rewards = []
    torch.manual_seed(seed)
    for it in range(iterations):
        scenes = make_scenarios(prompts, 0.3, seed=10_000 + 1000 * seed + it)
        style = instructions(prompts, 20_000 + 1000 * seed + it).repeat_interleave(
            group
        )
        drives = {k: v.repeat_interleave(group, 0) for k, v in scenes.items()}
        result = drive(drives, style, sampling(model), record=True)
        rewards = reward(result).view(prompts, group)
        advantages = group_advantages(rewards).reshape(-1, 1)
        n = prompts * group
        history = result["history"].reshape(n * STEPS, HISTORY, 3)
        styles = style[:, None].expand(n, STEPS).reshape(-1)
        taken = result["tokens"].reshape(-1, 1)

        def logp(net: MiniDriver) -> torch.Tensor:
            logits = net(history, styles)
            return torch.log_softmax(logits, -1).gather(1, taken).view(n, STEPS)

        with torch.no_grad():
            old_logp, ref_logp = logp(model), logp(reference)
        for _ in range(2):
            loss = grpo_loss(
                logp(model), old_logp, ref_logp, advantages, kl_coef=kl_coef
            )
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        history_of_rewards.append(rewards.mean().item())
    return history_of_rewards

def guarded(model: MiniDriver) -> Policy:
    """The model with the RSS filter between it and the actuators."""
    act = greedy(model)

    def policy(history, style, gap, speed, lead):
        accel, tokens = act(history, style, gap, speed, lead)
        safe, overridden = rss_filter(accel, gap, speed, lead, brake_min=BRAKE_MIN)
        return safe, tokens, overridden

    return policy