Imitation learning and action chunking

def expert_action(state: torch.Tensor) -> torch.Tensor:
    """The expert's steering: gentle near the lane centre, firm once off by 0.3 m.

    Near the centre it mostly damps the heading; beyond 0.3 m it steers back hard.
    Clean demonstrations rarely leave the gentle zone, so they never show the firm
    correction, which is the point of the toy.
    """
    e, psi = state[..., 0], state[..., 1]
    firm = 3.0 * torch.relu(e.abs() - 0.3) * torch.sign(e)
    return (-2.5 * psi - 0.05 * e - firm).clamp(-1.0, 1.0)


def step(
    state: torch.Tensor, action: torch.Tensor, disturbance: torch.Tensor
) -> torch.Tensor:
    """One time step: heading changes with the action plus a disturbance."""
    e, psi = state[..., 0], state[..., 1]
    e = e + DT * SPEED * torch.sin(psi)
    psi = psi + DT * action + disturbance
    return torch.stack([e, psi], dim=-1)

def collect(
    episodes: int,
    horizon: int,
    *,
    learner: nn.Module | None = None,
    action_noise: float = 0.0,
    disturbance: float = 0.01,
    generator: torch.Generator,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Rolls out and records (state, expert action) pairs.

    Who drives decides which states are seen; the label is always the expert's action:

    * behaviour cloning: the expert drives (``learner`` None, no noise);
    * DART: the expert drives with ``action_noise`` added to what is executed;
    * DAgger: the ``learner`` drives and the expert only labels.
    """
    state = torch.zeros(episodes, 2)
    state[:, 0] = 0.05 * torch.randn(episodes, generator=generator)
    states, labels = [], []
    for _ in range(horizon):
        label = expert_action(state)
        states.append(state)
        labels.append(label)
        if learner is not None:
            with torch.no_grad():
                executed = learner(state).squeeze(-1)
        else:
            executed = label + action_noise * torch.randn(episodes, generator=generator)
        noise = disturbance * torch.randn(episodes, generator=generator)
        state = step(state, executed, noise)
    return torch.cat(states), torch.cat(labels)

def behaviour_cloning(
    seed: int = 0, episodes: int = 20, horizon: int = 100
) -> nn.Module:
    """Fits the learner to clean expert demonstrations."""
    generator = torch.Generator().manual_seed(seed)
    torch.manual_seed(seed)
    learner = _policy()
    _fit(learner, *collect(episodes, horizon, generator=generator))
    return learner


def dart(
    seed: int = 0, episodes: int = 20, horizon: int = 100, noise: float = 0.5
) -> nn.Module:
    """Fits to demonstrations with perturbed executed actions; labels stay clean."""
    generator = torch.Generator().manual_seed(seed)
    torch.manual_seed(seed)
    learner = _policy()
    _fit(learner, *collect(episodes, horizon, action_noise=noise, generator=generator))
    return learner


def dagger(
    seed: int = 0, episodes: int = 20, horizon: int = 100, rounds: int = 4
) -> nn.Module:
    """Starts from behaviour cloning, then repeatedly lets the learner drive and the
    expert label the states it reaches, refitting on everything collected so far."""
    generator = torch.Generator().manual_seed(seed)
    torch.manual_seed(seed)
    learner = _policy()
    states, labels = collect(episodes, horizon, generator=generator)
    _fit(learner, states, labels)
    per_round = episodes // rounds
    for _ in range(rounds - 1):
        new_states, new_labels = collect(
            per_round, horizon, learner=learner, disturbance=0.03, generator=generator
        )
        states, labels = (
            torch.cat([states, new_states]),
            torch.cat([labels, new_labels]),
        )
        _fit(learner, states, labels)
    return learner

class RegressionPolicy(nn.Module):
    """Predicts the whole chunk with one forward pass and a mean-squared-error loss."""

    def __init__(self, hidden: int = 128) -> None:
        """Creates the MLP from an observation to a flattened chunk."""
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, hidden), nn.ReLU(), nn.Linear(hidden, CHUNK * 2)
        )

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        """Returns a chunk of shape (batch, CHUNK, 2)."""
        return self.net(obs).view(-1, CHUNK, 2)

    def loss(self, obs: torch.Tensor, chunk: torch.Tensor) -> torch.Tensor:
        """Mean squared error to the demonstrated chunk."""
        return ((self(obs) - chunk) ** 2).mean()

class CvaeChunkPolicy(nn.Module):
    """ACT-style policy: a style latent z explains which demonstration this was.

    During training an encoder reads the demonstrated chunk and produces z. The
    decoder maps (observation, z) to the chunk. At run time there is no
    demonstration to read, so z is drawn from the prior (or set to zero).
    """

    def __init__(self, hidden: int = 128) -> None:
        """Creates the encoder and decoder MLPs."""
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(CHUNK * 2, hidden), nn.ReLU(), nn.Linear(hidden, 2 * LATENT)
        )
        self.decoder = nn.Sequential(
            nn.Linear(1 + LATENT, hidden), nn.ReLU(), nn.Linear(hidden, CHUNK * 2)
        )

    def decode(self, obs: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        """Maps an observation and a latent to a chunk (batch, CHUNK, 2)."""
        return self.decoder(torch.cat([obs, z], dim=-1)).view(-1, CHUNK, 2)

    def loss(
        self, obs: torch.Tensor, chunk: torch.Tensor, beta: float = 0.05
    ) -> torch.Tensor:
        """Reconstruction error plus beta times the KL of the latent to N(0, I)."""
        mean, log_var = self.encoder(chunk.flatten(1)).chunk(2, dim=-1)
        z = mean + torch.exp(0.5 * log_var) * torch.randn_like(mean)
        recon = ((self.decode(obs, z) - chunk) ** 2).mean()
        kl = (-0.5 * (1 + log_var - mean**2 - log_var.exp()).sum(dim=-1)).mean()
        return recon + beta * kl

    @torch.no_grad()
    def act(self, obs: torch.Tensor, sample_prior: bool = True) -> torch.Tensor:
        """Predicts a chunk with z sampled from the prior, or z = 0 if not sampling."""
        shape = (len(obs), LATENT)
        z = torch.randn(shape) if sample_prior else torch.zeros(shape)
        return self.decode(obs, z)