"""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}")