Autoencoders and VAEs

"""The same task, done in different ways: why copying the average fails, and how a conditional VAE picks one way.

Run: python3 site/ml/foundations/examples/pytorch_vae.py

A robot starts near (0, 0) and must reach the far side of a 1 m table, past a round obstacle at (0, 0.5).
Each demonstration goes round it on the left or on the right. A policy sees two numbers, its start x and a
cue (−1 or +1, say a gap that is open on one side), and outputs a chunk of 10 waypoints: 20 numbers.
"""

import torch
from torch import nn

STEPS = torch.linspace(0.1, 1.0, 10)          # forward progress of the 10 waypoints
OBSTACLE, RADIUS = torch.tensor([0.0, 0.5]), 0.12


def demonstrations(count, generator, follow=0.5):
    """Observations (count, 2) and demonstrated chunks (count, 20). With probability `follow` the demonstrator
    goes round on the cue's side; otherwise on the other side. follow=0.5: the cue means nothing."""
    start = (torch.rand(count, generator=generator) - 0.5) * 0.2
    cue = torch.where(torch.rand(count, generator=generator) < 0.5, -1.0, 1.0)
    side = torch.where(torch.rand(count, generator=generator) < follow, cue, -cue)
    width = 0.25 + 0.1 * torch.rand(count, generator=generator)               # how wide this person swings
    x = start[:, None] * (1 - STEPS) + side[:, None] * width[:, None] * torch.sin(torch.pi * STEPS)
    chunk = torch.stack([x, STEPS.expand(count, 10)], dim=2).flatten(1)
    return torch.stack([start, cue], dim=1), chunk, side


def collides(chunks, observations):
    """Does the path from the start through the 10 waypoints come within the obstacle's radius anywhere?"""
    points = torch.cat([torch.stack([observations[:, 0], torch.zeros(len(chunks))], 1)[:, None], chunks.view(-1, 10, 2)], 1)
    s = torch.linspace(0, 1, 20)[None, None, :, None]
    along = (points[:, :-1, None] * (1 - s) + points[:, 1:, None] * s).reshape(len(chunks), -1, 2)
    return ((along - OBSTACLE).norm(dim=2) < RADIUS).any(dim=1)


def mlp(inputs, outputs):
    return nn.Sequential(nn.Linear(inputs, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, outputs))


class Plain(nn.Module):
    """Plain regression: observation in, chunk out, trained to match the demonstration."""
    def __init__(self):
        super().__init__()
        self.net = mlp(2, 20)

    def forward(self, observation, chunk=None):
        return self.net(observation), None, None


class CVAE(nn.Module):
    """A conditional VAE: the encoder reads observation and demonstrated chunk and says which 'style' z it was;
    the decoder turns observation and z back into a chunk."""
    def __init__(self, latent=1):
        super().__init__()
        self.latent = latent
        self.encoder = nn.Sequential(nn.Linear(2 + 20, 64), nn.ReLU(), nn.Linear(64, 2 * latent))
        self.decoder = mlp(2 + latent, 20)

    def forward(self, observation, chunk):
        mean, log_var = self.encoder(torch.cat([observation, chunk], 1)).chunk(2, dim=1)
        z = mean + torch.randn_like(mean) * (0.5 * log_var).exp()            # the reparameterisation trick
        return self.decoder(torch.cat([observation, z], 1)), mean, log_var

    def act(self, observation, z):
        """At test time there is no demonstration: z is chosen, then decoded."""
        return self.decoder(torch.cat([observation, z], 1))


def kl(mean, log_var):
    """How far each encoded distribution N(mean, var) is from N(0, 1), summed over z's numbers."""
    return (-0.5 * (1 + log_var - mean ** 2 - log_var.exp())).sum(dim=1).mean()


def train(kind, follow=0.5, beta=0.001, seed=0, epochs=200):          # beta chosen on validation starts
    torch.set_num_threads(1)
    torch.manual_seed(seed)
    generator = torch.Generator().manual_seed(seed)
    observations, chunks, _ = demonstrations(4000, generator, follow)
    policy = Plain() if kind == "plain" else CVAE()
    optimizer = torch.optim.Adam(policy.parameters(), lr=1e-3)
    for epoch in range(epochs):
        for batch in torch.randperm(4000, generator=generator).split(128):
            guess, mean, log_var = policy(observations[batch], chunks[batch])
            loss = ((guess - chunks[batch]) ** 2).mean()                      # reconstruction: match the demonstration
            if kind == "cvae":
                loss = loss + beta * kl(mean, log_var)                        # keep the styles near N(0, 1)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    return policy


def test(policy, follow, seed=9):
    """Collision rate over 1,000 new starts: plain, and for a CVAE with z = 0 and with z sampled from N(0, 1)."""
    generator = torch.Generator().manual_seed(seed)
    observations, chunks, side = demonstrations(1000, generator, follow)
    with torch.no_grad():
        if isinstance(policy, Plain):
            return {"plain": collides(policy(observations)[0], observations).float().mean().item()}
        zero = policy.act(observations, torch.zeros(1000, 1))
        sampled = policy.act(observations, torch.randn(1000, 1, generator=generator))
        mean, _ = policy.encoder(torch.cat([observations, chunks], 1)).chunk(2, dim=1)
    return {"z0": collides(zero, observations).float().mean().item(), "sampled": collides(sampled, observations).float().mean().item(),
            "z0_follows_cue": ((zero.view(-1, 10, 2)[:, 4, 0].sign()) == observations[:, 1]).float().mean().item(),
            "sampled_left": (sampled.view(-1, 10, 2)[:, 4, 0] < 0).float().mean().item(),
            "sampled_follows_cue": ((sampled.view(-1, 10, 2)[:, 4, 0].sign()) == observations[:, 1]).float().mean().item(),
            "z_left": mean[side < 0, 0].mean().item(), "z_right": mean[side > 0, 0].mean().item(),
            "z_with_cue": mean[side == observations[:, 1], 0].mean().item(), "z_against_cue": mean[side != observations[:, 1], 0].mean().item()}


if __name__ == "__main__":
    for follow in (0.5, 0.7):
        plain, cvae = train("plain", follow), train("cvae", follow)
        p, c = test(plain, follow), test(cvae, follow)
        print(f"cue followed {follow:.0%}: plain regression collides {p['plain']:.1%}; CVAE with z = 0 collides {c['z0']:.1%} "
              f"(follows the cue {c['z0_follows_cue']:.1%}), with z sampled {c['sampled']:.1%} ({c['sampled_left']:.0%} go left, "
              f"{c['sampled_follows_cue']:.0%} follow the cue); encoded z: left {c['z_left']:+.2f}, right {c['z_right']:+.2f}, "
              f"with the cue {c['z_with_cue']:+.2f}, against it {c['z_against_cue']:+.2f}")